fix unit test framework (#233)

Author: yiminc

internal_workflow_testsuite.go

@@ -770,6 +770,13 @@ func (w *workflowExecutorWrapper) Execute(ctx Context, input []byte) (result []b
 		}, false)
 	}
 
+	if !env.isChildWorkflow() {
+		// This is to prevent auto-forwarding mock clock before main workflow starts. For child workflow, we increase
+		// the counter in env.ExecuteChildWorkflow(). We cannot do it here for child workflow, because we need to make
+		// sure the counter is increased before returning from ExecuteChildWorkflow().
+		env.runningCount.Inc()
+	}
+
 	m := &mockWrapper{env: env, name: w.name, fn: w.fn, isWorkflow: true}
 	// This method is called by workflow's dispatcher. In this test suite, it is run in the main loop. We cannot block
 	// the main loop, but the mock could block if it is configured to wait. So we need to use a separate goroutinue to
@@ -790,12 +797,9 @@ func (w *workflowExecutorWrapper) Execute(ctx Context, input []byte) (result []b
 	// ExecuteUntilAllBlocked() so the main loop is not blocked. The dispatcher will unblock when getMockReturn() returns.
 	mockReadyChannel.Receive(ctx, &mockRet)
 
-	if env.isChildWorkflow() {
-		env.postCallback(func() {
-			// reduce runningCount after current workflow dispatcher run is blocked (aka ExecuteUntilAllBlocked() returns).
-			env.runningCount.Dec()
-		}, false)
-	}
+	// reduce runningCount to allow auto-forwarding mock clock after current workflow dispatcher run is blocked (aka
+	// ExecuteUntilAllBlocked() returns).
+	env.runningCount.Dec()
 
 	if mockRet != nil {
 		return m.executeMock(ctx, input, mockRet)

internal_workflow_testsuite_test.go

@@ -1045,36 +1045,29 @@ func (s *WorkflowTestSuiteUnitTest) Test_WorkflowFullyQualifiedName() {
 
 func (s *WorkflowTestSuiteUnitTest) Test_QueryWorkflow() {
 	queryType := "state"
-	stateWaitTimer, stateWaitActivity, stateDone := "wait for timer", "wait for activity", "done"
+	stateWaitSignal, stateWaitActivity, stateDone := "wait for signal", "wait for activity", "done"
 	workflowFn := func(ctx Context) error {
 		var state string
 		err := SetQueryHandler(ctx, queryType, func() (string, error) {
 			return state, nil
 		})
 		if err != nil {
-			state = err.Error()
 			return err
 		}
 
-		state = stateWaitTimer
-		err = NewTimer(ctx, time.Hour).Get(ctx, nil)
-		if err != nil {
-			state = err.Error()
-			return err
-		}
+		state = stateWaitSignal
+		var signalData string
+		GetSignalChannel(ctx, "query-signal").Receive(ctx, &signalData)
 
 		state = stateWaitActivity
 		ctx = WithActivityOptions(ctx, s.activityOptions)
 		err = ExecuteActivity(ctx, testActivityHello, "mock_delay").Get(ctx, nil)
 		if err != nil {
-			state = err.Error()
 			return err
 		}
-
 		state = stateDone
 		return err
 	}
-
 	RegisterWorkflow(workflowFn)
 
 	env := s.NewTestWorkflowEnvironment()
@@ -1086,9 +1079,10 @@ func (s *WorkflowTestSuiteUnitTest) Test_QueryWorkflow() {
 		s.NoError(err)
 		s.Equal(expected, state)
 	}
-	env.SetOnTimerScheduledListener(func(timerID string, duration time.Duration) {
-		verifyStateWithQuery(stateWaitTimer)
-	})
+	env.RegisterDelayedCallback(func() {
+		verifyStateWithQuery(stateWaitSignal)
+		env.SignalWorkflow("query-signal", "hello-query")
+	}, time.Hour)
 	env.OnActivity(testActivityHello, mock.Anything, mock.Anything).After(time.Hour).Return("hello_mock", nil)
 	env.SetOnActivityStartedListener(func(activityInfo *ActivityInfo, ctx context.Context, args EncodedValues) {
 		verifyStateWithQuery(stateWaitActivity)
@@ -1097,5 +1091,6 @@ func (s *WorkflowTestSuiteUnitTest) Test_QueryWorkflow() {
 
 	s.True(env.IsWorkflowCompleted())
 	s.NoError(env.GetWorkflowError())
+	env.AssertExpectations(s.T())
 	verifyStateWithQuery(stateDone)
 }