Minor: TimedFlowTests are sometimes flaky in CI since the flow result… (#3441)

Minor: TimedFlowTests are sometimes flaky in CI since the flow result future completes before the progress tracker is set to DONE. Instead of checking the progress tracker step immediately, wait for it with explicit timeout.

node/src/integration-test/kotlin/net/corda/node/services/TimedFlowTests.kt

@@ -11,6 +11,7 @@ import net.corda.core.identity.CordaX500Name
 import net.corda.core.identity.Party
 import net.corda.core.internal.FlowIORequest
 import net.corda.core.internal.ResolveTransactionsFlow
+import net.corda.core.internal.bufferUntilSubscribed
 import net.corda.core.internal.notary.NotaryServiceFlow
 import net.corda.core.internal.notary.TrustedAuthorityNotaryService
 import net.corda.core.internal.notary.UniquenessProvider
@@ -22,9 +23,9 @@ import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.utilities.ProgressTracker
 import net.corda.core.utilities.seconds
 import net.corda.node.internal.StartedNode
+import net.corda.node.services.config.FlowTimeoutConfiguration
 import net.corda.node.services.config.NodeConfiguration
 import net.corda.node.services.config.NotaryConfig
-import net.corda.node.services.config.FlowTimeoutConfiguration
 import net.corda.nodeapi.internal.DevIdentityGenerator
 import net.corda.nodeapi.internal.network.NetworkParametersCopier
 import net.corda.testing.common.internal.testNetworkParameters
@@ -43,8 +44,9 @@ import org.junit.BeforeClass
 import org.junit.Test
 import org.slf4j.MDC
 import java.security.PublicKey
+import java.util.concurrent.Future
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
-import kotlin.test.assertEquals
 import kotlin.test.assertNotEquals
 
 class TimedFlowTests {
@@ -136,13 +138,11 @@ class TimedFlowTests {
             val flow = NotaryFlow.Client(issueTx)
             val progressTracker = flow.progressTracker
             assertNotEquals(ProgressTracker.DONE, progressTracker.currentStep)
+            val progressTrackerDone = getDoneFuture(progressTracker)
+
             val notarySignatures = services.startFlow(flow).resultFuture.get()
             (issueTx + notarySignatures).verifyRequiredSignatures()
-            assertEquals(
+            progressTrackerDone.get()
-                    ProgressTracker.DONE,
-                    progressTracker.currentStep,
-                    "Ensure the same progress tracker object is re-used after flow restart"
-            )
         }
     }
 
@@ -155,14 +155,12 @@ class TimedFlowTests {
             }
             val flow = FinalityFlow(issueTx)
             val progressTracker = flow.progressTracker
+            assertNotEquals(ProgressTracker.DONE, progressTracker.currentStep)
+            val progressTrackerDone = getDoneFuture(flow.progressTracker)
 
             val stx = services.startFlow(flow).resultFuture.get()
             stx.verifyRequiredSignatures()
-            assertEquals(
+            progressTrackerDone.get()
-                    ProgressTracker.DONE,
-                    progressTracker.currentStep,
-                    "Ensure the same progress tracker object is re-used after flow restart"
-            )
         }
     }
 
@@ -175,6 +173,13 @@ class TimedFlowTests {
         )
     }
 
+    /** Returns a future that completes when the [progressTracker] reaches the [ProgressTracker.DONE] step. */
+    private fun getDoneFuture(progressTracker: ProgressTracker): Future<ProgressTracker.Change> {
+        return progressTracker.changes.takeFirst {
+            it.progressTracker.currentStep == ProgressTracker.DONE
+        }.timeout(5, TimeUnit.SECONDS).bufferUntilSubscribed().toBlocking().toFuture()
+    }
+
     @CordaService
     private class TestNotaryService(override val services: AppServiceHub, override val notaryIdentityKey: PublicKey) : TrustedAuthorityNotaryService() {
         override val uniquenessProvider = mock<UniquenessProvider>()