From 94d293b3b04531c17945bfc9018e1250af67a146 Mon Sep 17 00:00:00 2001
From: Vyacheslav Tamarin <slavabarsuk@ya.ru>
Date: Tue, 25 Jul 2023 17:04:09 +0300
Subject: [PATCH] Add generate-regression-suite param for cli

---
 utbot-cli-python/src/README.md                 |  8 ++++++++
 .../python/PythonGenerateTestsCommand.kt       | 18 +++++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/utbot-cli-python/src/README.md b/utbot-cli-python/src/README.md
index 7a963bfcc2..e469b0c1e6 100644
--- a/utbot-cli-python/src/README.md
+++ b/utbot-cli-python/src/README.md
@@ -77,6 +77,14 @@ Run generated tests:
 - `--test-framework [pytest|Unittest]`
 
   Test framework to be used.
+
+- `--runtime-exception-behaviour [PASS|FAIL]`
+
+  Expected behaviour for runtime exception.
+
+- `--generate-regression-suite`
+
+  Generate regression test suite. Default error test suite generation only.
   
 ### `run_python` options
 
diff --git a/utbot-cli-python/src/main/kotlin/org/utbot/cli/language/python/PythonGenerateTestsCommand.kt b/utbot-cli-python/src/main/kotlin/org/utbot/cli/language/python/PythonGenerateTestsCommand.kt
index 07af86255a..dcd0f8163c 100644
--- a/utbot-cli-python/src/main/kotlin/org/utbot/cli/language/python/PythonGenerateTestsCommand.kt
+++ b/utbot-cli-python/src/main/kotlin/org/utbot/cli/language/python/PythonGenerateTestsCommand.kt
@@ -9,8 +9,10 @@ import mu.KotlinLogging
 import org.parsers.python.PythonParser
 import org.utbot.framework.codegen.domain.RuntimeExceptionTestsBehaviour
 import org.utbot.framework.codegen.domain.TestFramework
+import org.utbot.framework.plugin.api.UtExecutionSuccess
 import org.utbot.python.PythonMethodHeader
 import org.utbot.python.PythonTestGenerationConfig
+import org.utbot.python.PythonTestSet
 import org.utbot.python.utils.RequirementsInstaller
 import org.utbot.python.TestFileInformation
 import org.utbot.python.code.PythonCode
@@ -102,6 +104,9 @@ class PythonGenerateTestsCommand : CliktCommand(
         .choice("PASS", "FAIL")
         .default("FAIL")
 
+    private val generateRegressionSuite by option("--generate-regression-suite", help = "Generate regression test suite")
+        .flag(default = false)
+
     private val testFramework: TestFramework
         get() =
             when (testFrameworkAsString) {
@@ -247,8 +252,19 @@ class PythonGenerateTestsCommand : CliktCommand(
         val (mypyStorage, _) = processor.sourceCodeAnalyze()
 
         logger.info("Generating tests...")
-        val testSets = processor.testGenerate(mypyStorage)
+        var testSets = processor.testGenerate(mypyStorage)
         if (testSets.isEmpty()) return
+        if (!generateRegressionSuite) {
+            testSets = testSets.map { testSet ->
+                PythonTestSet(
+                    testSet.method,
+                    testSet.executions.filter { it.result is UtExecutionSuccess },
+                    testSet.errors,
+                    testSet.mypyReport,
+                    testSet.classId
+                )
+            }
+        }
 
         logger.info("Saving tests...")
         val testCode = processor.testCodeGenerate(testSets)