fix: Fix race in TestEnvironment test causing flakes (#1266)

Author: mjameswh

packages/test/src/test-testenvironment.ts

@@ -2,7 +2,6 @@ import anyTest, { TestFn } from 'ava';
 import { v4 as uuid4 } from 'uuid';
 import { WorkflowFailedError } from '@temporalio/client';
 import { TestWorkflowEnvironment, workflowInterceptorModules } from '@temporalio/testing';
-import { Connection } from '@temporalio/testing/lib/connection';
 import { bundleWorkflowCode, WorkflowBundleWithSourceMap } from '@temporalio/worker';
 import {
   assertFromWorkflow,
@@ -79,11 +78,6 @@ test.serial('TestEnvironment can toggle between normal and skipped time', async
 });
 
 test.serial('TestEnvironment sleep can be used to delay activity completion', async (t) => {
-  // TODO: check why this fails on windows
-  if (process.platform === 'win32') {
-    t.pass();
-    return;
-  }
   const { client, nativeConnection, sleep } = t.context.testEnv;
 
   const worker = await Worker.create({
@@ -106,26 +100,14 @@ test.serial('TestEnvironment sleep can be used to delay activity completion', as
     t.is(winner, expectedWinner);
   };
   await worker.runUntil(async () => {
-    // TODO: there's an issue with the Java test server where if an activity
-    // does not complete before its scheduling workflow, time skipping stays
-    // locked.
-    // If the order of the below 2 statements is reversed, this test will hang.
     await run('activity');
     await run('timer');
   });
   t.pass();
 });
 
 test.serial('TestEnvironment sleep can be used to delay sending a signal', async (t) => {
-  // TODO: check why this fails on windows
-  if (process.platform === 'win32') {
-    t.pass();
-    return;
-  }
   const { client, nativeConnection, sleep } = t.context.testEnv;
-  // TODO: due to the test server issue mentioned in the test avove we need to manually unlock time skipping
-  // for the current test to balance out the time skipping lock counter.
-  await (t.context.testEnv.connection as Connection).testService.unlockTimeSkipping({});
 
   const worker = await Worker.create({
     connection: nativeConnection,

packages/test/src/workflows/testenv-test-workflows.ts

@@ -16,10 +16,18 @@ export async function raceActivityAndTimer(expectedWinner: 'timer' | 'activity')
   const timerShouldWin = expectedWinner === 'timer';
   const timerDuration = timerShouldWin ? 1_000_000 : 1_500_000;
   const activityDuration = timerShouldWin ? 1_500_000 : 1_000_000;
-  return await Promise.race([
-    sleep(timerDuration).then(() => 'timer'),
-    activities.sleep(activityDuration).then(() => 'activity'),
-  ]);
+
+  const timerPromise = sleep(timerDuration).then(() => 'timer');
+  const activityPromise = activities.sleep(activityDuration).then(() => 'activity');
+  const winner = await Promise.race([timerPromise, activityPromise]);
+
+  // TODO: there's an issue with the Java test server where if an activity does not complete
+  // before its scheduling workflow, time skipping stays locked, thus potentially causing errors
+  // on later tests. Work around this by making sure activity is completed before returning.
+  // See https://github.com/temporalio/sdk-java/issues/1138
+  await activityPromise;
+
+  return winner;
 }
 
 export async function waitOnSignalWithTimeout(): Promise<void> {

packages/worker/src/worker.ts

@@ -765,7 +765,7 @@ export class Worker {
    */
   shutdown(): void {
     if (this.state !== 'RUNNING') {
-      throw new IllegalStateError('Not running');
+      throw new IllegalStateError(`Not running. Current state: ${this.state}`);
     }
     this.state = 'STOPPING';
     this.nativeWorker