From 10411e3557b0fd670e97f1fcd74dcf3b20236a7d Mon Sep 17 00:00:00 2001
From: Ishan Taldekar <taldekar@amazon.com>
Date: Mon, 18 Nov 2024 19:34:53 -0500
Subject: [PATCH] Add extension to make multi-threaded calls run on a single
 thread for testing

---
 .../ThreadingUtilsStaticMockExtension.java    | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 plugin/tst/software/aws/toolkits/eclipse/amazonq/extensions/implementation/ThreadingUtilsStaticMockExtension.java

diff --git a/plugin/tst/software/aws/toolkits/eclipse/amazonq/extensions/implementation/ThreadingUtilsStaticMockExtension.java b/plugin/tst/software/aws/toolkits/eclipse/amazonq/extensions/implementation/ThreadingUtilsStaticMockExtension.java
new file mode 100644
index 000000000..82716c0bd
--- /dev/null
+++ b/plugin/tst/software/aws/toolkits/eclipse/amazonq/extensions/implementation/ThreadingUtilsStaticMockExtension.java
@@ -0,0 +1,51 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.eclipse.amazonq.extensions.implementation;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mockStatic;
+
+import org.junit.jupiter.api.extension.AfterAllCallback;
+import org.junit.jupiter.api.extension.BeforeAllCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.mockito.MockedStatic;
+
+import software.aws.toolkits.eclipse.amazonq.extensions.api.StaticMockExtension;
+import software.aws.toolkits.eclipse.amazonq.util.ThreadingUtils;
+
+public final class ThreadingUtilsStaticMockExtension extends StaticMockExtension<ThreadingUtils>
+        implements BeforeAllCallback, AfterAllCallback {
+
+    private MockedStatic<ThreadingUtils> threadingUtilsStaticMock = null;
+
+    @Override
+    public MockedStatic<ThreadingUtils> getStaticMock() {
+        return threadingUtilsStaticMock;
+    }
+
+    @Override
+    public void beforeAll(final ExtensionContext context) {
+        threadingUtilsStaticMock = mockStatic(ThreadingUtils.class);
+
+        threadingUtilsStaticMock.when(() -> ThreadingUtils.executeAsyncTask(any(Runnable.class))).then(invocation -> {
+            Runnable task = invocation.getArgument(0);
+            task.run();
+            return null;
+        });
+        threadingUtilsStaticMock.when(() -> ThreadingUtils.executeAsyncTaskAndReturnFuture(any(Runnable.class)))
+                .then(invocation -> {
+                    Runnable task = invocation.getArgument(0);
+                    task.run();
+                    return null;
+                });
+    }
+
+    @Override
+    public void afterAll(final ExtensionContext context) throws Exception {
+        if (threadingUtilsStaticMock != null) {
+            threadingUtilsStaticMock.close();
+        }
+    }
+
+}