Fix CI test flake (#675)

Author: Sushisource

.github/workflows/per-pr.yml

@@ -65,7 +65,7 @@ jobs:
       - uses: actions-rs/cargo@v1
         with:
           command: test
-          args: -- --include-ignored
+          args: -- --include-ignored --nocapture
       - uses: actions/upload-artifact@v3
         with:
           name: coverage-report

core/src/core_tests/workflow_tasks.rs

@@ -2112,15 +2112,12 @@ async fn continue_as_new_preserves_some_values() {
     wes_attrs.search_attributes = Some(search);
     wes_attrs.retry_policy = Some(retry_policy);
     let mut mock_client = mock_workflow_client();
-    let hist = {
+    let t = {
         let mut t = TestHistoryBuilder::default();
         t.add(wes_attrs.clone());
         t.add_full_wf_task();
         t
     };
-    mock_client.expect_poll_workflow_task().returning(move |_| {
-        Ok(hist_to_poll_resp(&hist, wfid.to_owned(), ResponseType::AllHistory).resp)
-    });
     mock_client
         .expect_complete_workflow_task()
         .returning(move |mut c| {
@@ -2135,8 +2132,8 @@ async fn continue_as_new_preserves_some_values() {
             }
             Ok(Default::default())
         });
-
-    let worker = Worker::new_test(test_worker_cfg().build().unwrap(), mock_client);
+    let mock = single_hist_mock_sg(wfid, t, vec![ResponseType::AllHistory], mock_client, true);
+    let worker = mock_worker(mock);
     let r = worker.poll_workflow_activation().await.unwrap();
     worker
         .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
@@ -2149,6 +2146,7 @@ async fn continue_as_new_preserves_some_values() {
         ))
         .await
         .unwrap();
+    worker.shutdown().await;
 }
 
 #[rstest]
@@ -2737,15 +2735,12 @@ async fn use_compatible_version_flag(
 ) {
     let wfid = "fake_wf_id";
     let mut mock_client = mock_workflow_client();
-    let hist = {
+    let t = {
         let mut t = TestHistoryBuilder::default();
         t.add_by_type(EventType::WorkflowExecutionStarted);
         t.add_full_wf_task();
         t
     };
-    mock_client.expect_poll_workflow_task().returning(move |_| {
-        Ok(hist_to_poll_resp(&hist, wfid.to_owned(), ResponseType::AllHistory).resp)
-    });
     let compat_flag_expected = match intent {
         VersioningIntent::Unspecified => !different_tq,
         VersioningIntent::Compatible => true,
@@ -2770,7 +2765,8 @@ async fn use_compatible_version_flag(
             Ok(Default::default())
         });
 
-    let worker = Worker::new_test(test_worker_cfg().build().unwrap(), mock_client);
+    let mock = single_hist_mock_sg(wfid, t, vec![ResponseType::AllHistory], mock_client, true);
+    let worker = mock_worker(mock);
     let r = worker.poll_workflow_activation().await.unwrap();
     let task_queue = if different_tq {
         "enchi cat!".to_string()
@@ -2806,6 +2802,7 @@ async fn use_compatible_version_flag(
         .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(r.run_id, cmd))
         .await
         .unwrap();
+    worker.shutdown().await;
 }
 
 #[tokio::test]

core/src/worker/activities/activity_heartbeat_manager.rs

@@ -450,8 +450,7 @@ mod test {
         hm.shutdown().await;
     }
 
-    /// Ensure that heartbeat can be called from a tight loop without any throttle_interval, resulting in two
-    /// interactions with the server - one immediately and one after 500ms after the throttle_interval.
+    /// Ensure that heartbeat can be called from a tight loop and correctly throttle
     #[tokio::test]
     async fn process_tight_loop_and_shutdown() {
         let mut mock_client = mock_workflow_client();

core/src/worker/activities/local_activities.rs

@@ -1246,7 +1246,7 @@ mod tests {
         // Wait more than the timeout before grabbing the task
         sleep(timeout + Duration::from_millis(10)).await;
 
-        assert!(dbg!(lam.next_pending().await.unwrap()).is_timeout(false));
+        assert!(lam.next_pending().await.unwrap().is_timeout(false));
         assert_eq!(lam.num_in_backoff(), 0);
         assert_eq!(lam.num_outstanding(), 0);
     }

core/src/worker/workflow/machines/local_activity_state_machine.rs

@@ -1250,7 +1250,6 @@ mod tests {
         });
 
         let mut worker = build_fake_sdk(mock_cfg);
-        dbg!("Past thing");
         worker.register_wf(DEFAULT_WORKFLOW_TYPE, la_wf);
         worker.register_activity(
             DEFAULT_ACTIVITY_TYPE,

core/src/worker/workflow/machines/upsert_search_attributes_state_machine.rs

@@ -208,7 +208,9 @@ mod tests {
     use temporal_sdk_core_api::Worker;
     use temporal_sdk_core_protos::{
         coresdk::{
-            workflow_commands::SetPatchMarker, workflow_completion::WorkflowActivationCompletion,
+            workflow_activation::{workflow_activation_job, WorkflowActivationJob},
+            workflow_commands::SetPatchMarker,
+            workflow_completion::WorkflowActivationCompletion,
             AsJsonPayloadExt,
         },
         temporal::api::{
@@ -332,24 +334,38 @@ mod tests {
             t.add_has_change_marker(&patch_id, false);
         }
         t.add_upsert_search_attrs_for_patch(&[patch_id.clone()]);
+        t.add_we_signaled("hi", vec![]);
         t.add_full_wf_task();
         t.add_workflow_execution_completed();
 
         let mut mp = MockPollCfg::from_resp_batches(
             "fakeid",
             t,
-            [ResponseType::ToTaskNum(1), ResponseType::AllHistory],
+            [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)],
             mock_workflow_client(),
         );
         // Ensure the upsert command has an empty map when not using the patched command
         if !with_patched_cmd {
             mp.completion_asserts = Some(Box::new(|wftc| {
-                assert_matches!(wftc.commands.first().and_then(|c| c.attributes.as_ref()).unwrap(),
-            Attributes::UpsertWorkflowSearchAttributesCommandAttributes(attrs)
-            if attrs.search_attributes.as_ref().unwrap().indexed_fields.is_empty())
+                let cmd_attrs = wftc
+                    .commands
+                    .first()
+                    .and_then(|c| c.attributes.as_ref())
+                    .unwrap();
+                if matches!(
+                    cmd_attrs,
+                    Attributes::CompleteWorkflowExecutionCommandAttributes(_)
+                ) {
+                    return;
+                }
+                assert_matches!(cmd_attrs,
+                    Attributes::UpsertWorkflowSearchAttributesCommandAttributes(attrs)
+                    if attrs.search_attributes.clone().unwrap_or_default().indexed_fields.is_empty());
             }));
         }
-        let core = mock_worker(build_mock_pollers(mp));
+        let mut mock = build_mock_pollers(mp);
+        mock.worker_cfg(|w| w.max_cached_workflows = 1);
+        let core = mock_worker(mock);
 
         let mut ver_upsert = HashMap::new();
         ver_upsert.insert(
@@ -379,6 +395,12 @@ mod tests {
         .unwrap();
         // Now ensure that encountering the upsert in history works fine
         let act = core.poll_workflow_activation().await.unwrap();
+        assert_matches!(
+            act.jobs.as_slice(),
+            [WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::SignalWorkflow(_)),
+            }]
+        );
         core.complete_execution(&act.run_id).await;
     }
 }