Update time management of python test generation

utbot-python/src/main/kotlin/org/utbot/python/PythonEngine.kt

@@ -31,7 +31,6 @@ class PythonEngine(
     private val pythonPath: String,
     private val fuzzedConcreteValues: List<PythonFuzzedConcreteValue>,
     private val timeoutForRun: Long,
-    private val initialCoveredLines: Set<Int>,
     private val pythonTypeStorage: PythonTypeStorage,
 ) {
 
@@ -141,17 +140,20 @@ class PythonEngine(
 
     fun fuzzing(parameters: List<Type>, isCancelled: () -> Boolean, until: Long): Flow<FuzzingExecutionFeedback> = flow {
         val additionalModules = parameters.flatMap { it.pythonModules() }
-        val coveredLines = initialCoveredLines.toMutableSet()
 
         ServerSocket(0).use { serverSocket ->
             logger.info { "Server port: ${serverSocket.localPort}" }
-            val manager = PythonWorkerManager(
-                serverSocket,
-                pythonPath,
-                until,
-                { constructEvaluationInput(it) },
-                timeoutForRun.toInt()
-            )
+            val manager = try {
+                PythonWorkerManager(
+                    serverSocket,
+                    pythonPath,
+                    until,
+                    { constructEvaluationInput(it) },
+                    timeoutForRun.toInt()
+                )
+            } catch (_: TimeoutException) {
+                return@flow
+            }
             logger.info { "Executor manager was created successfully" }
 
             fun fuzzingResultHandler(
@@ -195,7 +197,6 @@ class PythonEngine(
 
                         is PythonEvaluationSuccess -> {
                             val coveredInstructions = evaluationResult.coverage.coveredInstructions
-                            coveredInstructions.forEach { coveredLines.add(it.lineNumber) }
 
                             val summary = arguments
                                 .zip(methodUnderTest.arguments)
@@ -277,7 +278,6 @@ class PythonEngine(
                     cache.add(pair, result)
                     emit(result.fuzzingExecutionFeedback)
                     return@PythonFuzzing result.fuzzingPlatformFeedback
-
                 }.fuzz(pmd)
             }
         }

utbot-python/src/main/kotlin/org/utbot/python/PythonTestCaseGenerator.kt

@@ -22,15 +22,14 @@ import org.utbot.python.newtyping.mypy.MypyReportLine
 import org.utbot.python.newtyping.mypy.getErrorNumber
 import org.utbot.python.newtyping.utils.getOffsetLine
 import org.utbot.python.typing.MypyAnnotations
+import org.utbot.python.utils.ExecutionWithTimeoutMode
+import org.utbot.python.utils.GenerationLimitManager
 import org.utbot.python.utils.PriorityCartesianProduct
+import org.utbot.python.utils.TimeoutMode
 import java.io.File
 
 private val logger = KotlinLogging.logger {}
 
-private const val COVERAGE_LIMIT = 150
-private const val ADDITIONAL_LIMIT = 5
-private const val INVALID_EXECUTION_LIMIT = 10
-
 class PythonTestCaseGenerator(
     private val withMinimization: Boolean = true,
     private val directoriesForSysPath: Set<String>,
@@ -120,6 +119,10 @@ class PythonTestCaseGenerator(
 
     fun generate(method: PythonMethod, until: Long): PythonTestSet {
         storageForMypyMessages.clear()
+        val limitManager = GenerationLimitManager(
+            ExecutionWithTimeoutMode,
+            until,
+        )
 
         val typeStorage = PythonTypeStorage.get(mypyStorage)
 
@@ -135,10 +138,6 @@ class PythonTestCaseGenerator(
         val coveredLines = mutableSetOf<Int>()
         var generated = 0
 
-        var additionalLimit = ADDITIONAL_LIMIT
-        val typeInferenceCancellation =
-            { isCancelled() || System.currentTimeMillis() >= until || additionalLimit <= 0 }
-
         logger.info("Start test generation for ${method.name}")
         substituteTypeParameters(method, typeStorage).forEach { newMethod ->
             inferAnnotations(
@@ -148,7 +147,7 @@ class PythonTestCaseGenerator(
                 hintCollector,
                 mypyReportLine,
                 mypyConfigFile,
-                typeInferenceCancellation
+                limitManager,
             ) { functionType ->
                 val args = (functionType as FunctionType).arguments
 
@@ -161,22 +160,12 @@ class PythonTestCaseGenerator(
                     pythonPath,
                     constants,
                     timeoutForRun,
-                    coveredLines,
                     PythonTypeStorage.get(mypyStorage)
                 )
 
-                var invalidExecutionLimit = INVALID_EXECUTION_LIMIT
-                var coverageLimit = COVERAGE_LIMIT
-                var coveredBefore = coveredLines.size
-
                 var feedback: InferredTypeFeedback = SuccessFeedback
 
-                val fuzzerCancellation = {
-                        typeInferenceCancellation()
-                                || coverageLimit == 0
-                                || additionalLimit == 0
-                                || invalidExecutionLimit == 0
-                }
+                val fuzzerCancellation = { isCancelled() || limitManager.isCancelled() }
                 val startTime = System.currentTimeMillis()
 
                 engine.fuzzing(args, fuzzerCancellation, until).collect {
@@ -186,30 +175,26 @@ class PythonTestCaseGenerator(
                             executions += it.utFuzzedExecution
                             missingLines = updateCoverage(it.utFuzzedExecution, coveredLines, missingLines)
                             feedback = SuccessFeedback
+                            limitManager.addSuccessExecution()
                         }
                         is InvalidExecution -> {
                             errors += it.utError
                             feedback = SuccessFeedback
                         }
                         is ArgumentsTypeErrorFeedback -> {
-                            invalidExecutionLimit--
                             feedback = InvalidTypeFeedback
+                            limitManager.addInvalidExecution()
                         }
                         is TypeErrorFeedback -> {
-                            invalidExecutionLimit--
                             feedback = InvalidTypeFeedback
+                            limitManager.addInvalidExecution()
                         }
                     }
-                    if (missingLines?.size == 0) {
-                        additionalLimit--
-                    }
-                    val coveredAfter = coveredLines.size
-                    if (coveredAfter == coveredBefore) {
-                        coverageLimit--
-                    }
+                    limitManager.missedLines = missingLines?.size
+
                     logger.info { "Time ${System.currentTimeMillis() - startTime}: $generated, $missingLines" }
-                    coveredBefore = coveredAfter
                 }
+                limitManager.restart()
                 feedback
             }
         }
@@ -250,7 +235,7 @@ class PythonTestCaseGenerator(
         hintCollector: HintCollector,
         report: List<MypyReportLine>,
         mypyConfigFile: File,
-        isCancelled: () -> Boolean,
+        limitManager: GenerationLimitManager,
         annotationHandler: suspend (Type) -> InferredTypeFeedback,
     ) {
         val namesInModule = mypyStorage.names
@@ -259,6 +244,7 @@ class PythonTestCaseGenerator(
             .filter {
                 it.length < 4 || !it.startsWith("__") || !it.endsWith("__")
             }
+        val typeInferenceCancellation = { isCancelled() || limitManager.isCancelled() }
 
         val algo = BaselineAlgorithm(
             typeStorage,
@@ -277,14 +263,15 @@ class PythonTestCaseGenerator(
         )
 
         runBlocking breaking@{
-            if (isCancelled()) {
+            if (typeInferenceCancellation()) {
                 return@breaking
             }
 
-            algo.run(hintCollector.result, isCancelled, annotationHandler)
+            val iterationNumber = algo.run(hintCollector.result, typeInferenceCancellation, annotationHandler)
 
-            val existsAnnotation = method.definition.type
-            if (existsAnnotation.arguments.all { it.pythonTypeName() != "typing.Any" }) {
+            if (iterationNumber == 1) {
+                limitManager.mode = TimeoutMode
+                val existsAnnotation = method.definition.type
                 annotationHandler(existsAnnotation)
             }
         }

utbot-python/src/main/kotlin/org/utbot/python/evaluation/PythonWorkerManager.kt

@@ -24,8 +24,8 @@ class PythonWorkerManager(
 
     var timeout: Long = 0
     lateinit var process: Process
-    lateinit var workerSocket: Socket
-    lateinit var codeExecutor: PythonCodeExecutor
+    private lateinit var workerSocket: Socket
+    private lateinit var codeExecutor: PythonCodeExecutor
 
     init {
         connect()
@@ -39,7 +39,7 @@ class PythonWorkerManager(
             "localhost",
             serverSocket.localPort.toString(),
             "--logfile", logfile.absolutePath,
-            //"--loglevel", "DEBUG",
+            "--loglevel", "INFO",  // "DEBUG", "INFO", "ERROR"
         ))
         timeout = max(until - processStartTime, 0)
         workerSocket = try {

utbot-python/src/main/kotlin/org/utbot/python/newtyping/inference/TypeInferenceAPI.kt

@@ -8,7 +8,7 @@ abstract class TypeInferenceAlgorithm {
         hintCollectorResult: HintCollectorResult,
         isCancelled: () -> Boolean,
         annotationHandler: suspend (Type) -> InferredTypeFeedback,
-    )
+    ): Int
 }
 
 sealed interface InferredTypeFeedback

utbot-python/src/main/kotlin/org/utbot/python/newtyping/inference/baseline/BaselineAlgorithm.kt

@@ -39,20 +39,26 @@ class BaselineAlgorithm(
         hintCollectorResult: HintCollectorResult,
         isCancelled: () -> Boolean,
         annotationHandler: suspend (Type) -> InferredTypeFeedback,
-    ) {
+    ): Int {
         val generalRating = createGeneralTypeRating(hintCollectorResult, storage)
         val initialState = getInitialState(hintCollectorResult, generalRating)
         val states: MutableList<BaselineAlgorithmState> = mutableListOf(initialState)
         val fileForMypyRuns = TemporaryFileManager.assignTemporaryFile(tag = "mypy.py")
+        var iterationCounter = 0
 
         run breaking@ {
             while (states.isNotEmpty()) {
                 if (isCancelled())
                     return@breaking
                 logger.debug("State number: ${states.size}")
+                iterationCounter++
+
                 val state = chooseState(states)
                 val newState = expandState(state, storage)
                 if (newState != null) {
+                    if (iterationCounter == 1) {
+                        annotationHandler(initialState.signature)
+                    }
                     logger.info("Checking ${newState.signature.pythonTypeRepresentation()}")
                     if (checkSignature(newState.signature as FunctionType, fileForMypyRuns, configFile)) {
                         logger.debug("Found new state!")
@@ -71,6 +77,7 @@ class BaselineAlgorithm(
                 }
             }
         }
+        return iterationCounter
     }
 
     private fun checkSignature(signature: FunctionType, fileForMypyRuns: File, configFile: File): Boolean {

utbot-python/src/main/kotlin/org/utbot/python/utils/GenerationLimitManager.kt

@@ -0,0 +1,75 @@
+package org.utbot.python.utils
+
+import kotlin.math.min
+
+class GenerationLimitManager(
+    // global settings
+    var mode: LimitManagerMode,
+    val until: Long,
+
+    // local settings: one type inference iteration
+    var executions: Int = 150,
+    var invalidExecutions: Int = 10,
+    var additionalExecutions: Int = 5,
+    var missedLines: Int? = null,
+) {
+    private val initExecution = executions
+    private val initInvalidExecutions = invalidExecutions
+    private val initAdditionalExecutions = additionalExecutions
+    private val initMissedLines = missedLines
+
+    fun restart() {
+        executions = initExecution
+        invalidExecutions = initInvalidExecutions
+        additionalExecutions = initAdditionalExecutions
+        missedLines = initMissedLines
+    }
+
+    fun addSuccessExecution() {
+        executions -= 1
+    }
+    fun addInvalidExecution() {
+        invalidExecutions -= 1
+    }
+
+    fun isCancelled(): Boolean {
+        return mode.isCancelled(this)
+    }
+}
+
+interface LimitManagerMode {
+    fun isCancelled(manager: GenerationLimitManager): Boolean
+}
+
+object MaxCoverageMode : LimitManagerMode {
+    override fun isCancelled(manager: GenerationLimitManager): Boolean {
+        return manager.missedLines?.equals(0) == true
+    }
+}
+
+object TimeoutMode : LimitManagerMode {
+    override fun isCancelled(manager: GenerationLimitManager): Boolean {
+        return System.currentTimeMillis() >= manager.until
+    }
+}
+
+object ExecutionMode : LimitManagerMode {
+    override fun isCancelled(manager: GenerationLimitManager): Boolean {
+        if (manager.invalidExecutions <= 0 || manager.executions <= 0) {
+            return min(manager.invalidExecutions, 0) + min(manager.executions, 0) <= manager.additionalExecutions
+        }
+        return false
+    }
+}
+
+object MaxCoverageWithTimeoutMode : LimitManagerMode {
+    override fun isCancelled(manager: GenerationLimitManager): Boolean {
+        return MaxCoverageMode.isCancelled(manager) || TimeoutMode.isCancelled(manager)
+    }
+}
+
+object ExecutionWithTimeoutMode : LimitManagerMode {
+    override fun isCancelled(manager: GenerationLimitManager): Boolean {
+        return ExecutionMode.isCancelled(manager) || TimeoutMode.isCancelled(manager)
+    }
+}