Fix calculation of expected test count

core/src/main/kotlin/com/malinskiy/marathon/execution/progress/ProgressReporter.kt

@@ -50,8 +50,8 @@ class ProgressReporter(private val configuration: Configuration) {
     fun aggregateResult(): Boolean =
         reporters.isNotEmpty() && reporters.values.all { it.aggregateResult() }
 
-    fun testCountExpectation(poolId: DevicePoolId, size: Int) {
-        execute(poolId) { it.testCountExpectation(size) }
+    fun addTests(poolId: DevicePoolId, count: Int) {
+        execute(poolId) { it.addTests(count) }
     }
 
     fun removeTests(poolId: DevicePoolId, count: Int) {

core/src/main/kotlin/com/malinskiy/marathon/execution/progress/tracker/PoolProgressTracker.kt

@@ -106,14 +106,12 @@ class PoolProgressTracker(private val strictMode: Boolean) {
         }
     }
 
-    fun testCountExpectation(size: Int) {
-        expectedTestCount.set(size)
+    fun addTests(count: Int) {
+        expectedTestCount.updateAndGet { it + count }
     }
 
     fun removeTests(count: Int) {
-        expectedTestCount.updateAndGet {
-            it - count
-        }
+        expectedTestCount.updateAndGet { it - count }
     }
 
     fun progress(): Float =

core/src/main/kotlin/com/malinskiy/marathon/execution/queue/QueueActor.kt

@@ -68,7 +68,7 @@ class QueueActor(
                 testResultReporter.addShard(msg.shard)
                 val testsToAdd = msg.shard.tests + msg.shard.flakyTests
                 queue.addAll(testsToAdd)
-                progressReporter.addRetries(poolId, testsToAdd.size)
+                progressReporter.addTests(poolId, testsToAdd.size)
                 flakyTests = flakyTests + msg.shard.flakyTests
 
                 if (queue.isNotEmpty()) {

core/src/test/kotlin/com/malinskiy/marathon/execution/progress/ProgressReporterTest.kt

@@ -21,7 +21,7 @@ class ProgressReporterTest {
         val test2 = MarathonTest("com.example", "SimpleTest", "method2", emptyList(), TestComponentInfo())
         val test3 = MarathonTest("com.example", "SimpleTest", "method3", emptyList(), TestComponentInfo())
 
-        reporter.testCountExpectation(poolId, 3)
+        reporter.addTests(poolId, 3)
         reporter.progress().shouldBeEqualTo(.0f)
 
         /**
@@ -76,6 +76,43 @@ class ProgressReporterTest {
         progress.shouldBeEqualTo(6 / 6f)
     }
 
+    @Test
+    fun shouldReportProgressForOnePoolWithMultipleShards() {
+        val poolId = DevicePoolId("testpool")
+
+        val test1 = MarathonTest("com.example", "SimpleTest", "method1", emptyList(), TestComponentInfo())
+        val test2 = MarathonTest("com.example", "SimpleTest", "method2", emptyList(), TestComponentInfo())
+        val test3 = MarathonTest("com.example", "SimpleTest", "method3", emptyList(), TestComponentInfo())
+
+        // Add the first shard
+        reporter.addTests(poolId, 2)
+        reporter.progress().shouldBeEqualTo(.0f)
+
+        /**
+         * test 1 passed
+         */
+        reporter.testStarted(poolId, deviceInfo, test1)
+        reporter.testPassed(poolId, deviceInfo, test1)
+        reporter.progress().shouldBeEqualTo(1 / 2f)
+
+        // Add the second shard
+        reporter.addTests(poolId, 1)
+
+        /**
+         * test 2 passed
+         */
+        reporter.testStarted(poolId, deviceInfo, test2)
+        reporter.testPassed(poolId, deviceInfo, test2)
+        reporter.progress().shouldBeEqualTo(2 / 3f)
+
+        /**
+         * test 3 passed
+         */
+        reporter.testStarted(poolId, deviceInfo, test3)
+        reporter.testPassed(poolId, deviceInfo, test3)
+        reporter.progress().shouldBeEqualTo(3 / 3f)
+    }
+
     @Test
     fun shouldReportProgressForOnePoolWithRuntimeDiscovery() {
         val poolId = DevicePoolId("testpool")
@@ -84,7 +121,7 @@ class ProgressReporterTest {
         val test1 = MarathonTest("com.example", "SimpleTest", "method[1]", emptyList(), TestComponentInfo())
         val test2 = MarathonTest("com.example", "SimpleTest", "method[2]", emptyList(), TestComponentInfo())
 
-        reporter.testCountExpectation(poolId, 1)
+        reporter.addTests(poolId, 1)
         reporter.progress().shouldBeEqualTo(.0f)
 
         /**

core/src/test/kotlin/com/malinskiy/marathon/execution/progress/tracker/PoolProgressTrackerTest.kt

@@ -17,7 +17,7 @@ class PoolProgressTrackerTest {
     @Test
     fun nonStrictMode_case1() {
         val tracker = PoolProgressTracker(strictMode = false)
-        tracker.testCountExpectation(1)
+        tracker.addTests(1)
         tracker.testStarted(test)
         tracker.testPassed(test)
         tracker.testFailed(test)
@@ -29,7 +29,7 @@ class PoolProgressTrackerTest {
     @Test
     fun strictMode_case1() {
         val tracker = PoolProgressTracker(strictMode = true)
-        tracker.testCountExpectation(1)
+        tracker.addTests(1)
         tracker.testStarted(test)
         tracker.testPassed(test)
         tracker.testFailed(test)
@@ -41,7 +41,7 @@ class PoolProgressTrackerTest {
     @Test
     fun all_incomplete() {
         val tracker = PoolProgressTracker(strictMode = false).apply {
-            testCountExpectation(1)
+            addTests(1)
         }
         tracker.aggregateResult() shouldBeEqualTo false
     }
@@ -50,7 +50,7 @@ class PoolProgressTrackerTest {
     fun withRetries() {
         val tracker = PoolProgressTracker(strictMode = false)
 
-        tracker.testCountExpectation(1)
+        tracker.addTests(1)
         tracker.testStarted(test)
         tracker.testFailed(test)
         tracker.addTestRetries(1)
@@ -77,7 +77,7 @@ class PoolProgressTrackerTest {
             componentInfo = TestComponentInfo()
         )
 
-        tracker.testCountExpectation(1)
+        tracker.addTests(1)
         tracker.testStarted(test0)
         tracker.testPassed(test0)
         tracker.testStarted(test1)