From 7150aae84f4ba7263f248b005574d276f6cd48de Mon Sep 17 00:00:00 2001
From: Andrew Baldwin <abald.cg.ide@aldogroup.com>
Date: Fri, 19 Jan 2024 08:44:05 -0500
Subject: [PATCH] Fix hanging test

---
 .github/workflows/tests.yml                   |  2 +-
 .../tests/SwarmUserClassPicker.test.tsx       | 36 ++++++++++++-------
 locust/webui/src/test/setup.ts                |  6 +++-
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index e0fd12af7d..a400d3345c 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -47,7 +47,7 @@ jobs:
       - run: python -m pip install tox
       - run: python -m tox -e ${{ matrix.tox }}
 
-  lint_typecheck_webui:
+  lint_typecheck_test_webui:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
diff --git a/locust/webui/src/components/SwarmForm/tests/SwarmUserClassPicker.test.tsx b/locust/webui/src/components/SwarmForm/tests/SwarmUserClassPicker.test.tsx
index d2ebeed5ce..1101b6d0ca 100644
--- a/locust/webui/src/components/SwarmForm/tests/SwarmUserClassPicker.test.tsx
+++ b/locust/webui/src/components/SwarmForm/tests/SwarmUserClassPicker.test.tsx
@@ -1,4 +1,4 @@
-import { act, fireEvent, waitFor } from '@testing-library/react';
+import { act, fireEvent } from '@testing-library/react';
 import { http } from 'msw';
 import { setupServer } from 'msw/node';
 import { afterEach, afterAll, beforeAll, describe, test, expect, vi } from 'vitest';
@@ -47,7 +47,7 @@ const defaultProps = {
   setSwarm: updateSwarm,
 };
 
-describe('SwarmForm', () => {
+describe('SwarmUserClassPicker', () => {
   beforeAll(() => server.listen());
   afterEach(() => {
     server.resetHandlers();
@@ -95,6 +95,8 @@ describe('SwarmForm', () => {
       },
     );
 
+    vi.useFakeTimers();
+
     act(() => {
       // Open modal
       fireEvent.click(getAllByRole('button')[0]);
@@ -104,12 +106,15 @@ describe('SwarmForm', () => {
       fireEvent.click(getByRole('button', { name: 'Save' }));
     });
 
-    await waitFor(async () => {
-      const submittedData = updateUserSettings.mock.calls[0][0];
-      if (submittedData) {
-        expect(submittedData).toEqual({ ...mockUsers.Example, userClassName: 'Example' });
-      }
+    await act(async () => {
+      await vi.runAllTimersAsync();
     });
+
+    expect(updateUserSettings).toHaveBeenCalled();
+
+    const submittedData = updateUserSettings.mock.calls[0][0];
+
+    expect(submittedData).toEqual({ ...mockUsers.Example, userClassName: 'Example' });
   });
 
   test('should allow for configuring tasks, host, fixed_count, and weight to be modified', async () => {
@@ -130,6 +135,8 @@ describe('SwarmForm', () => {
       },
     );
 
+    vi.useFakeTimers();
+
     act(() => {
       // Open modal
       fireEvent.click(getAllByRole('button')[0]);
@@ -152,12 +159,15 @@ describe('SwarmForm', () => {
       fireEvent.click(getByRole('button', { name: 'Save' }));
     });
 
-    await waitFor(async () => {
-      const submittedData = updateUserSettings.mock.calls[0][0];
-      if (submittedData) {
-        expect(submittedData).toEqual({ ...updatedUser, userClassName: 'Example' });
-        expect(store.getState().swarm.users.Example).toEqual(updatedUser);
-      }
+    await act(async () => {
+      await vi.runAllTimersAsync();
     });
+
+    expect(updateUserSettings).toHaveBeenCalled();
+
+    const submittedData = updateUserSettings.mock.calls[0][0];
+
+    expect(submittedData).toEqual({ ...updatedUser, userClassName: 'Example' });
+    expect(store.getState().swarm.users.Example).toEqual(updatedUser);
   });
 });
diff --git a/locust/webui/src/test/setup.ts b/locust/webui/src/test/setup.ts
index f9e29a8cac..48f361429b 100644
--- a/locust/webui/src/test/setup.ts
+++ b/locust/webui/src/test/setup.ts
@@ -1,5 +1,5 @@
 import { cleanup } from '@testing-library/react';
-import { afterEach, vi } from 'vitest';
+import { afterAll, afterEach, vi } from 'vitest';
 
 import { TEST_BASE_API } from 'test/constants';
 import { swarmStateMock } from 'test/mocks/swarmState.mock';
@@ -25,3 +25,7 @@ vi.mock('echarts', async () => {
 
   return { ...actual, init: (...args: any[]) => actual.init(...args, { width: 100, height: 100 }) };
 });
+
+afterAll(() => {
+  vi.clearAllMocks();
+});