Fix some test parsing issues

JetBrains-Research · Aug 29, 2023 · b9c82a1 · b9c82a1
1 parent 836a32d
commit b9c82a1
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/main/kotlin/org/jetbrains/research/testspark/tools/llm/generation/TestsAssembler.kt b/src/main/kotlin/org/jetbrains/research/testspark/tools/llm/generation/TestsAssembler.kt
@@ -100,7 +100,8 @@ class TestsAssembler(
     fun returnTestSuite(packageName: String): TestSuiteGeneratedByLLM? {
         try {
             val testSuite = TestSuiteGeneratedByLLM(project)
-            rawText = rawText.split("```")[1]
+            if(rawText.contains("```"))
+                rawText = rawText.split("```")[1]
 
             testSuite.packageString = packageName
 
@@ -170,7 +171,9 @@ class TestsAssembler(
 
                 // Get test's body
                 // remove opening bracket
-                var testBody: String = rawTest.split("{")[1].trim()
+                var testBody =rawTest.split("{").toMutableList().apply {
+                    removeFirst()
+                }.joinToString("{").trim()
 
                 // remove closing bracket