Add a test for Noop delay interrupt

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Jul 20, 2024 · de9de39 · de9de39
1 parent 3958bc1
commit de9de39
Showing 1 changed file with 46 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/flowframework/workflow/NoOpStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/NoOpStepTests.java
@@ -15,6 +15,9 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.opensearch.flowframework.common.CommonValue.DELAY_FIELD;
 
@@ -50,6 +53,49 @@ public void testNoOpStepDelay() throws IOException, InterruptedException {
         assertTrue(System.nanoTime() - start > 900_000_000L);
     }
 
+    public void testNoOpStepInterrupt() throws IOException, InterruptedException {
+        NoOpStep noopStep = new NoOpStep();
+        WorkflowData delayData = new WorkflowData(Map.of(DELAY_FIELD, "5s"), null, null);
+
+        CountDownLatch latch = new CountDownLatch(1);
+        // Fetch errors from the separate thread
+        AtomicReference<AssertionError> assertionError = new AtomicReference<>();
+
+        Thread testThread = new Thread(() -> {
+            try {
+                PlainActionFuture<WorkflowData> future = noopStep.execute(
+                    "nodeId",
+                    delayData,
+                    Collections.emptyMap(),
+                    Collections.emptyMap(),
+                    Collections.emptyMap()
+                );
+                try {
+                    future.actionGet();
+                } catch (Exception e) {
+                    // Ignore the IllegalStateExcption/InterruptedExcpetion
+                }
+                assertTrue(future.isDone());
+                assertTrue(future.isCancelled());
+                assertTrue(Thread.currentThread().isInterrupted());
+            } catch (AssertionError e) {
+                assertionError.set(e);
+            } finally {
+                latch.countDown();
+            }
+        });
+
+        testThread.start();
+        Thread.sleep(100);
+        testThread.interrupt();
+
+        latch.await(1, TimeUnit.SECONDS);
+
+        if (assertionError.get() != null) {
+            throw assertionError.get();
+        }
+    }
+
     public void testNoOpStepParse() throws IOException {
         NoOpStep noopStep = new NoOpStep();
         WorkflowData delayData = new WorkflowData(Map.of(DELAY_FIELD, "foo"), null, null);