Merge pull request #950 from MarathonLabs/fix/device-actor-return-batch

fix(core): improve return batch behavior
MarathonLabs · Jun 25, 2024 · c4dd149 · c4dd149
2 parents 6ada7f1 + d913336
commit c4dd149
Show file tree

Hide file tree

Showing 5 changed files with 86 additions and 9 deletions.
diff --git a/core/src/main/kotlin/com/malinskiy/marathon/execution/DevicePoolActor.kt b/core/src/main/kotlin/com/malinskiy/marathon/execution/DevicePoolActor.kt
@@ -86,6 +86,7 @@ class DevicePoolActor(
     }
 
     private suspend fun deviceReturnedTestBatch(device: Device, batch: TestBatch, reason: String) {
+        logger.debug { "pool $poolId: device ${device.serialNumber} returned test batch" }
         queue.send(QueueMessage.ReturnBatch(device.toDeviceInfo(), batch, reason))
     }
 
@@ -98,7 +99,7 @@ class DevicePoolActor(
     }
 
     // Requests a batch of tests for a random device from the list of devices not running tests at the moment.
-    // When @avoidingDevice is not null, attemtps to send the request for any other device whenever available.
+    // When @avoidingDevice is not null, attempts to send the request for any other device whenever available.
     private suspend fun maybeRequestBatch(avoidingDevice: Device? = null) {
         val availableDevices = devices.values.asSequence()
             .map { it as DeviceActor }

diff --git a/core/src/main/kotlin/com/malinskiy/marathon/execution/DevicePoolMessage.kt b/core/src/main/kotlin/com/malinskiy/marathon/execution/DevicePoolMessage.kt
@@ -11,10 +11,10 @@ sealed class DevicePoolMessage {
         object Terminate : FromScheduler()
     }
 
-    sealed class FromDevice(val device: Device) : DevicePoolMessage() {
-        class IsReady(device: Device) : FromDevice(device)
-        class CompletedTestBatch(device: Device, val results: TestBatchResults) : FromDevice(device)
-        class ReturnTestBatch(device: Device, val batch: TestBatch, val reason: String) : FromDevice(device)
+    sealed class FromDevice(open val device: Device) : DevicePoolMessage() {
+        data class IsReady(override val device: Device) : FromDevice(device)
+        data class CompletedTestBatch(override val device: Device, val results: TestBatchResults) : FromDevice(device)
+        data class ReturnTestBatch(override val device: Device, val batch: TestBatch, val reason: String) : FromDevice(device)
     }
 
     sealed class FromQueue : DevicePoolMessage() {

diff --git a/core/src/main/kotlin/com/malinskiy/marathon/execution/device/DeviceActor.kt b/core/src/main/kotlin/com/malinskiy/marathon/execution/device/DeviceActor.kt
@@ -17,10 +17,12 @@ import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CompletionHandler
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.async
 import kotlinx.coroutines.channels.SendChannel
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import kotlin.coroutines.CoroutineContext
 
 class DeviceActor(
@@ -216,7 +218,9 @@ class DeviceActor(
 
     private fun returnBatchAnd(batch: TestBatch, reason: String, completionHandler: CompletionHandler = {}): Job {
         return launch {
-            pool.send(DevicePoolMessage.FromDevice.ReturnTestBatch(device, batch, reason))
+            withContext(NonCancellable) {
+                pool.send(DevicePoolMessage.FromDevice.ReturnTestBatch(device, batch, reason))
+            }
         }.apply {
             invokeOnCompletion(completionHandler)
         }

diff --git a/core/src/main/kotlin/com/malinskiy/marathon/execution/queue/QueueActor.kt b/core/src/main/kotlin/com/malinskiy/marathon/execution/queue/QueueActor.kt
@@ -73,7 +73,7 @@ class QueueActor(
             }
 
             is QueueMessage.ReturnBatch -> {
-                onReturnBatch(msg.device, msg.batch)
+                onReturnBatch(msg.device, msg.batch, msg.reason)
             }
         }
     }
@@ -124,8 +124,8 @@ class QueueActor(
         }
     }
 
-    private suspend fun onReturnBatch(device: DeviceInfo, batch: TestBatch) {
-        logger.debug { "onReturnBatch ${device.serialNumber}" }
+    private suspend fun onReturnBatch(device: DeviceInfo, batch: TestBatch, reason: String) {
+        logger.debug { "onReturnBatch ${device.serialNumber}. reason=$reason" }
 
         val uncompletedTests = batch.tests
         val results = uncompletedTests.map {

diff --git a/core/src/test/kotlin/com/malinskiy/marathon/execution/device/DeviceActorTest.kt b/core/src/test/kotlin/com/malinskiy/marathon/execution/device/DeviceActorTest.kt
@@ -0,0 +1,72 @@
+package com.malinskiy.marathon.execution.device
+
+import com.malinskiy.marathon.config.Configuration
+import com.malinskiy.marathon.config.strategy.BatchingStrategyConfiguration
+import com.malinskiy.marathon.config.vendor.VendorConfiguration
+import com.malinskiy.marathon.device.DevicePoolId
+import com.malinskiy.marathon.execution.DevicePoolMessage
+import com.malinskiy.marathon.execution.TestStatus
+import com.malinskiy.marathon.generateTest
+import com.malinskiy.marathon.test.StubDevice
+import com.malinskiy.marathon.test.TestBatch
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.runBlocking
+import org.amshove.kluent.shouldBeEqualTo
+import org.junit.jupiter.api.AfterEach
+import org.junit.jupiter.api.BeforeEach
+import org.junit.jupiter.api.Test
+import java.io.File
+
+class DeviceActorTest {
+    lateinit var job: Job
+    lateinit var poolChannel: Channel<DevicePoolMessage>
+
+    @BeforeEach
+    fun setup() {
+        poolChannel = Channel()
+        job = Job()
+    }
+
+    @AfterEach
+    fun teardown() {
+        job.cancel()
+    }
+
+    @Test
+    fun `terminate while in progress`() {
+        val devicePoolId = DevicePoolId("test")
+        val device = StubDevice(prepareTimeMillis = 1000L, testTimeMillis = 10000L)
+        val actor = DeviceActor(
+            devicePoolId, poolChannel, defaultConfiguration.copy(
+                uncompletedTestRetryQuota = 0,
+                batchingStrategy = BatchingStrategyConfiguration.FixedSizeBatchingStrategyConfiguration(size = 1)
+            ), device, job, Dispatchers.Unconfined
+        )
+
+        runBlocking {
+            val test1 = generateTest()
+            val testBatch = TestBatch(listOf(test1))
+            device.executionResults = mapOf(test1 to Array(1) { TestStatus.FAILURE })
+
+            actor.send(DeviceEvent.Initialize)
+            var message = poolChannel.receive()
+            message.shouldBeEqualTo(DevicePoolMessage.FromDevice.IsReady(device))
+
+            actor.send(DeviceEvent.Execute(testBatch))
+            actor.send(DeviceEvent.Terminate)
+
+            message = poolChannel.receive()
+            message.shouldBeEqualTo(DevicePoolMessage.FromDevice.ReturnTestBatch(device, testBatch, "Device serial-1 terminated"))
+        }
+    }
+
+    private val defaultConfiguration = Configuration.Builder(
+        name = "",
+        outputDir = File(""),
+    ).apply {
+        vendorConfiguration = VendorConfiguration.StubVendorConfiguration
+        analyticsTracking = false
+    }.build()
+}