NIFI-16183 Add active-polling stopConnector(Duration) to connector mock framework

nifi-connector-mock-bundle/nifi-connector-mock-api/src/main/java/org/apache/nifi/mock/connector/server/ConnectorTestRunner.java

@@ -32,6 +32,7 @@
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeoutException;
 
 public interface ConnectorTestRunner extends Closeable {
 
@@ -144,6 +145,19 @@ public interface ConnectorTestRunner extends Closeable {
      */
     void stopConnector();
 
+    /**
+     * Stops the Connector, waiting up to the given timeout for it to reach a stopped state. Implementations
+     * should actively poll the Connector's state until it is stopped or the timeout elapses, which is more
+     * tolerant of a flow that takes a while to quiesce (for example a failed table still draining) than the
+     * default stop budget.
+     *
+     * @param timeout the maximum duration to wait for the Connector to stop
+     * @throws TimeoutException if the timeout elapses before the Connector stops
+     */
+    default void stopConnector(final Duration timeout) throws TimeoutException {
+        stopConnector();
+    }
+
     /**
      * Blocks until the Connector has received at least one FlowFile, or until the specified timeout elapses.
      *

nifi-connector-mock-bundle/nifi-connector-mock-server/src/main/java/org/apache/nifi/mock/connector/server/StandardConnectorMockServer.java

@@ -102,7 +102,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.jar.JarFile;
 import java.util.stream.Stream;
 
@@ -113,6 +113,8 @@ public class StandardConnectorMockServer implements ConnectorMockServer {
     private static final String NAR_DEPENDENCIES_PATH = "NAR-INF/bundled-dependencies";
     private static final String CONNECTOR_WAR_MANIFEST_PATH = "META-INF/nifi-connector";
     private static final String WAR_EXTENSION = ".war";
+    private static final Duration DEFAULT_STOP_TIMEOUT = Duration.ofSeconds(60);
+    private static final long STOP_POLL_INTERVAL_MILLIS = 250L;
 
     private Bundle systemBundle;
     private Set<Bundle> bundles;
@@ -362,13 +364,35 @@ public void startConnector() {
 
     @Override
     public void stopConnector() {
+        // The no-arg convenience method keeps its unchecked contract: a timeout at the default budget is not
+        // something a caller is expected to recover from, so wrap the checked TimeoutException.
         try {
-            connectorNode.stop(flowEngine).get(10, TimeUnit.SECONDS);
-        } catch (final InterruptedException e) {
-            Thread.currentThread().interrupt();
-            throw new RuntimeException("Interrupted while waiting for connector to stop", e);
-        } catch (final Exception e) {
-            throw new RuntimeException("Failed to stop Connector", e);
+            stopConnector(DEFAULT_STOP_TIMEOUT);
+        } catch (final TimeoutException e) {
+            throw new RuntimeException(e.getMessage(), e);
+        }
+    }
+
+    @Override
+    public void stopConnector(final Duration maxWaitTime) throws TimeoutException {
+        // Initiate the asynchronous stop, then actively poll the Connector's state until it reports STOPPED.
+        // The node flips its state to STOPPED at the same point it completes the stop future, and it retries a
+        // failed component stop internally (every 10 seconds), so polling the state rides through those retries
+        // up to maxWaitTime instead of being capped by a single fixed blocking wait.
+        connectorNode.stop(flowEngine);
+
+        final long expirationTime = System.currentTimeMillis() + maxWaitTime.toMillis();
+        while (connectorNode.getCurrentState() != ConnectorState.STOPPED) {
+            if (System.currentTimeMillis() > expirationTime) {
+                throw new TimeoutException("Timed out waiting for the Connector to stop after " + maxWaitTime);
+            }
+
+            try {
+                Thread.sleep(STOP_POLL_INTERVAL_MILLIS);
+            } catch (final InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new RuntimeException("Interrupted while waiting for the Connector to stop", e);
+            }
         }
     }
 

nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java

@@ -28,10 +28,12 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -65,6 +67,25 @@ public void testCreateStartAndStopGenerateAndUpdateConnector() throws IOExceptio
         }
     }
 
+    @Test
+    public void testStopConnectorWithTimeoutStopsRunningConnector() throws IOException {
+        try (final ConnectorTestRunner testRunner = new StandardConnectorTestRunner.Builder()
+                .connectorClassName("org.apache.nifi.mock.connectors.GenerateAndLog")
+                .narLibraryDirectory(new File("target/libDir"))
+                .build()) {
+
+            testRunner.startConnector();
+
+            // Exercises the timeout-aware overload: it initiates the asynchronous stop and actively polls the
+            // Connector's state until it reaches STOPPED, returning as soon as it does rather than after a single
+            // fixed blocking wait. Because start is asynchronous, a stop issued immediately afterwards may have to
+            // ride through the node's internal stop retries (every 10 seconds) before the state settles, so the
+            // budget is generous enough to stay deterministic on a slow CI runner; the poll still returns the
+            // instant the Connector reports STOPPED.
+            assertDoesNotThrow(() -> testRunner.stopConnector(Duration.ofSeconds(120)));
+        }
+    }
+
     @Test
     public void testConnectorWithMissingBundleFailsValidate() throws IOException {
 

nifi-connector-mock-bundle/nifi-connector-mock/src/main/java/org/apache/nifi/mock/connector/StandardConnectorTestRunner.java

@@ -50,6 +50,7 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.TimeoutException;
 
 public class StandardConnectorTestRunner implements ConnectorTestRunner, Closeable {
     private final File narLibraryDirectory;
@@ -198,6 +199,11 @@ public void stopConnector() {
         mockServer.stopConnector();
     }
 
+    @Override
+    public void stopConnector(final Duration timeout) throws TimeoutException {
+        mockServer.stopConnector(timeout);
+    }
+
     @Override
     public void waitForDataIngested(final Duration maxWaitTime) {
         mockServer.waitForDataIngested(maxWaitTime);