@temporalio/core-bridge 0.16.4 → 0.18.0

package/sdk-core/core/src/core_tests/local_activities.rs

@@ -0,0 +1,328 @@
+use crate::{
+    replay::{default_wes_attribs, TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE},
+    test_help::{build_mock_pollers, mock_core, MockPollCfg, ResponseType, TEST_Q},
+    Core,
+};
+use anyhow::anyhow;
+use futures::future::join_all;
+use std::{
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+use temporal_client::mocks::mock_gateway;
+use temporal_sdk::{LocalActivityOptions, TestRustWorker, WfContext, WorkflowResult};
+use temporal_sdk_core_protos::{
+    coresdk::{common::RetryPolicy, AsJsonPayloadExt},
+    temporal::api::{enums::v1::EventType, failure::v1::Failure},
+};
+use tokio::sync::Barrier;
+
+async fn echo(e: String) -> anyhow::Result<String> {
+    Ok(e)
+}
+
+/// This test verifies that when replaying we are able to resolve local activities whose data we
+/// don't see until after the workflow issues the command
+#[rstest::rstest]
+#[case::replay(true, true)]
+#[case::not_replay(false, true)]
+#[case::replay_cache_off(true, false)]
+#[case::not_replay_cache_off(false, false)]
+#[tokio::test]
+async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached: bool) {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    t.add_full_wf_task();
+    t.add_local_activity_result_marker(1, "1", b"echo".into());
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_full_wf_task();
+    t.add_workflow_execution_completed();
+
+    let wf_id = "fakeid";
+    let mock = mock_gateway();
+    let resps = if replay {
+        vec![ResponseType::AllHistory]
+    } else {
+        vec![1.into(), 2.into(), ResponseType::AllHistory]
+    };
+    let mh = MockPollCfg::from_resp_batches(wf_id, t, resps, mock);
+    let mut mock = build_mock_pollers(mh);
+    if cached {
+        mock.worker_cfg(TEST_Q, |cfg| cfg.max_cached_workflows = 1);
+    }
+    let core = mock_core(mock);
+    let mut worker = TestRustWorker::new(Arc::new(core), TEST_Q.to_string(), None);
+
+    worker.register_wf(
+        DEFAULT_WORKFLOW_TYPE.to_owned(),
+        |ctx: WfContext| async move {
+            let la = ctx.local_activity(LocalActivityOptions {
+                activity_type: "echo".to_string(),
+                input: "hi".as_json_payload().expect("serializes fine"),
+                ..Default::default()
+            });
+            ctx.timer(Duration::from_secs(1)).await;
+            la.await;
+            Ok(().into())
+        },
+    );
+    worker.register_activity("echo", echo);
+    worker
+        .submit_wf(wf_id.to_owned(), DEFAULT_WORKFLOW_TYPE.to_owned(), vec![])
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}
+
+pub async fn local_act_fanout_wf(ctx: WfContext) -> WorkflowResult<()> {
+    let las: Vec<_> = (1..=50)
+        .map(|i| {
+            ctx.local_activity(LocalActivityOptions {
+                activity_type: "echo".to_string(),
+                input: format!("Hi {}", i)
+                    .as_json_payload()
+                    .expect("serializes fine"),
+                ..Default::default()
+            })
+        })
+        .collect();
+    ctx.timer(Duration::from_secs(1)).await;
+    join_all(las).await;
+    Ok(().into())
+}
+
+#[tokio::test]
+async fn local_act_many_concurrent() {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    t.add_full_wf_task();
+    for i in 1..=50 {
+        t.add_local_activity_result_marker(i, &i.to_string(), b"echo".into());
+    }
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_full_wf_task();
+    t.add_workflow_execution_completed();
+
+    let wf_id = "fakeid";
+    let mock = mock_gateway();
+    let mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 3], mock);
+    let mock = build_mock_pollers(mh);
+    let core = mock_core(mock);
+    let mut worker = TestRustWorker::new(Arc::new(core), TEST_Q.to_string(), None);
+
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE.to_owned(), local_act_fanout_wf);
+    worker.register_activity("echo", |str: String| async move { Ok(str) });
+    worker
+        .submit_wf(wf_id.to_owned(), DEFAULT_WORKFLOW_TYPE.to_owned(), vec![])
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}
+
+/// Verifies that local activities which take more than a workflow task timeout will cause
+/// us to issue additional (empty) WFT completions with the force flag on, thus preventing timeout
+/// of WFT while the local activity continues to execute.
+///
+/// The test with shutdown verifies if we call shutdown while the local activity is running that
+/// shutdown does not complete until it's finished.
+#[rstest::rstest]
+#[case::with_shutdown(true)]
+#[case::normal_complete(false)]
+#[tokio::test]
+async fn local_act_heartbeat(#[case] shutdown_middle: bool) {
+    let mut t = TestHistoryBuilder::default();
+    let wft_timeout = Duration::from_millis(200);
+    let mut wes_short_wft_timeout = default_wes_attribs();
+    wes_short_wft_timeout.workflow_task_timeout = Some(wft_timeout.into());
+    t.add(
+        EventType::WorkflowExecutionStarted,
+        wes_short_wft_timeout.into(),
+    );
+    t.add_full_wf_task();
+    // Task created by WFT heartbeat
+    t.add_full_wf_task();
+    t.add_workflow_task_scheduled_and_started();
+
+    let wf_id = "fakeid";
+    let mock = mock_gateway();
+    // Allow returning incomplete history more than once, as the wft timeout can be timing sensitive
+    // and might poll an extra time
+    let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 2, 2, 2], mock);
+    mh.enforce_correct_number_of_polls = false;
+    let mut mock = build_mock_pollers(mh);
+    mock.worker_cfg(TEST_Q, |wc| wc.max_cached_workflows = 1);
+    let core = Arc::new(mock_core(mock));
+    let mut worker = TestRustWorker::new(core.clone(), TEST_Q.to_string(), None);
+    let shutdown_barr: &'static Barrier = Box::leak(Box::new(Barrier::new(2)));
+
+    worker.register_wf(
+        DEFAULT_WORKFLOW_TYPE.to_owned(),
+        |ctx: WfContext| async move {
+            ctx.local_activity(LocalActivityOptions {
+                activity_type: "echo".to_string(),
+                input: "hi".as_json_payload().expect("serializes fine"),
+                ..Default::default()
+            })
+            .await;
+            Ok(().into())
+        },
+    );
+    worker.register_activity("echo", move |str: String| async move {
+        if shutdown_middle {
+            shutdown_barr.wait().await;
+        }
+        // Take slightly more than two workflow tasks
+        tokio::time::sleep(wft_timeout.mul_f32(2.2)).await;
+        Ok(str)
+    });
+    worker
+        .submit_wf(wf_id.to_owned(), DEFAULT_WORKFLOW_TYPE.to_owned(), vec![])
+        .await
+        .unwrap();
+    let (_, runres) = tokio::join!(
+        async {
+            if shutdown_middle {
+                shutdown_barr.wait().await;
+                core.shutdown().await;
+            }
+        },
+        worker.run_until_done()
+    );
+    runres.unwrap();
+}
+
+#[rstest::rstest]
+#[case::retry_then_pass(true)]
+#[case::retry_until_fail(false)]
+#[tokio::test]
+async fn local_act_fail_and_retry(#[case] eventually_pass: bool) {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_workflow_task_scheduled_and_started();
+
+    let wf_id = "fakeid";
+    let mock = mock_gateway();
+    let mh = MockPollCfg::from_resp_batches(wf_id, t, [1], mock);
+    let mock = build_mock_pollers(mh);
+    let core = mock_core(mock);
+    let mut worker = TestRustWorker::new(Arc::new(core), TEST_Q.to_string(), None);
+
+    worker.register_wf(
+        DEFAULT_WORKFLOW_TYPE.to_owned(),
+        move |ctx: WfContext| async move {
+            let la_res = ctx
+                .local_activity(LocalActivityOptions {
+                    activity_type: "echo".to_string(),
+                    input: "hi".as_json_payload().expect("serializes fine"),
+                    retry_policy: RetryPolicy {
+                        initial_interval: Some(Duration::from_millis(50).into()),
+                        backoff_coefficient: 1.2,
+                        maximum_interval: None,
+                        maximum_attempts: 5,
+                        non_retryable_error_types: vec![],
+                    },
+                    ..Default::default()
+                })
+                .await;
+            if eventually_pass {
+                assert!(la_res.completed_ok())
+            } else {
+                assert!(la_res.failed())
+            }
+            Ok(().into())
+        },
+    );
+    let attempts: &'static _ = Box::leak(Box::new(AtomicUsize::new(0)));
+    worker.register_activity("echo", move |_: String| async move {
+        // Succeed on 3rd attempt (which is ==2 since fetch_add returns prev val)
+        if 2 == attempts.fetch_add(1, Ordering::Relaxed) && eventually_pass {
+            Ok(())
+        } else {
+            Err(anyhow!("Oh no I failed!"))
+        }
+    });
+    worker
+        .submit_wf(wf_id.to_owned(), DEFAULT_WORKFLOW_TYPE.to_owned(), vec![])
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+    let expected_attempts = if eventually_pass { 3 } else { 5 };
+    assert_eq!(expected_attempts, attempts.load(Ordering::Relaxed));
+}
+
+#[tokio::test]
+async fn local_act_retry_long_backoff_uses_timer() {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    t.add_local_activity_fail_marker(
+        1,
+        "1",
+        Failure::application_failure("la failed".to_string(), false),
+    );
+    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_full_wf_task();
+    t.add_local_activity_fail_marker(
+        2,
+        "2",
+        Failure::application_failure("la failed".to_string(), false),
+    );
+    let timer_started_event_id = t.add_get_event_id(EventType::TimerStarted, None);
+    t.add_timer_fired(timer_started_event_id, "2".to_string());
+    t.add_full_wf_task();
+    t.add_workflow_execution_completed();
+
+    let wf_id = "fakeid";
+    let mock = mock_gateway();
+    let mh = MockPollCfg::from_resp_batches(
+        wf_id,
+        t,
+        [1.into(), 2.into(), ResponseType::AllHistory],
+        mock,
+    );
+    let mut mock = build_mock_pollers(mh);
+    mock.worker_cfg(TEST_Q, |w| w.max_cached_workflows = 1);
+    let core = mock_core(mock);
+    let mut worker = TestRustWorker::new(Arc::new(core), TEST_Q.to_string(), None);
+
+    worker.register_wf(
+        DEFAULT_WORKFLOW_TYPE.to_owned(),
+        |ctx: WfContext| async move {
+            let la_res = ctx
+                .local_activity(LocalActivityOptions {
+                    activity_type: "echo".to_string(),
+                    input: "hi".as_json_payload().expect("serializes fine"),
+                    retry_policy: RetryPolicy {
+                        initial_interval: Some(Duration::from_millis(65).into()),
+                        // This will make the second backoff 65 seconds, plenty to use timer
+                        backoff_coefficient: 1_000.,
+                        maximum_interval: Some(Duration::from_secs(600).into()),
+                        maximum_attempts: 3,
+                        non_retryable_error_types: vec![],
+                    },
+                    ..Default::default()
+                })
+                .await;
+            assert!(la_res.failed());
+            // Extra timer just to have an extra workflow task which we can return full history for
+            ctx.timer(Duration::from_secs(1)).await;
+            Ok(().into())
+        },
+    );
+    worker.register_activity("echo", move |_: String| async move {
+        Result::<(), _>::Err(anyhow!("Oh no I failed!"))
+    });
+    worker
+        .submit_wf(wf_id.to_owned(), DEFAULT_WORKFLOW_TYPE.to_owned(), vec![])
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}

package/sdk-core/{src → core/src}/core_tests/mod.rs

@@ -1,25 +1,23 @@
 mod activity_tasks;
 mod child_workflows;
 mod determinism;
+mod local_activities;
 mod queries;
 mod replay_flag;
-mod retry;
 mod workers;
 mod workflow_cancels;
 mod workflow_tasks;
 
 use crate::{
     errors::{PollActivityError, PollWfError},
-    pollers::MockManualGateway,
-    test_help::{
-        build_fake_core, canned_histories, fake_sg_opts, hist_to_poll_resp, ResponseType, TEST_Q,
-    },
+    test_help::{build_fake_core, canned_histories, hist_to_poll_resp, ResponseType, TEST_Q},
     Core, CoreInitOptionsBuilder, CoreSDK, WorkerConfigBuilder,
 };
 use futures::FutureExt;
 use std::time::Duration;
+use temporal_client::mocks::{fake_sg_opts, mock_manual_gateway};
 use temporal_sdk_core_protos::{
-    coresdk::workflow_completion::WfActivationCompletion,
+    coresdk::workflow_completion::WorkflowActivationCompletion,
     temporal::api::workflowservice::v1::PollActivityTaskQueueResponse,
 };
 use tokio::{sync::Barrier, time::sleep};
@@ -30,7 +28,7 @@ async fn after_shutdown_server_is_not_polled() {
     let core = build_fake_core("fake_wf_id", t, &[1]);
     let res = core.poll_workflow_activation(TEST_Q).await.unwrap();
     assert_eq!(res.jobs.len(), 1);
-    core.complete_workflow_activation(WfActivationCompletion::empty(TEST_Q, res.run_id))
+    core.complete_workflow_activation(WorkflowActivationCompletion::empty(TEST_Q, res.run_id))
         .await
         .unwrap();
     core.shutdown().await;
@@ -46,7 +44,7 @@ lazy_static::lazy_static! {
 }
 #[tokio::test]
 async fn shutdown_interrupts_both_polls() {
-    let mut mock_gateway = MockManualGateway::new();
+    let mut mock_gateway = mock_manual_gateway();
     mock_gateway
         .expect_poll_activity_task()
         .times(1)
@@ -96,7 +94,6 @@ async fn shutdown_interrupts_both_polls() {
             .build()
             .unwrap(),
     )
-    .await
     .unwrap();
     tokio::join! {
         async {

package/sdk-core/{src → core/src}/core_tests/queries.rs

@@ -1,16 +1,20 @@
 use crate::{
-    pollers::MockServerGatewayApis,
     test_help::{
         canned_histories, hist_to_poll_resp, mock_core, MocksHolder, ResponseType, TEST_Q,
     },
     Core,
 };
-use std::collections::{HashMap, VecDeque};
+use std::{
+    collections::{HashMap, VecDeque},
+    time::Duration,
+};
+use temporal_client::mocks::mock_gateway;
+
 use temporal_sdk_core_protos::{
     coresdk::{
-        workflow_activation::{wf_activation_job, WfActivationJob},
-        workflow_commands::{CompleteWorkflowExecution, QueryResult, QuerySuccess, StartTimer},
-        workflow_completion::WfActivationCompletion,
+        workflow_activation::{workflow_activation_job, WorkflowActivationJob},
+        workflow_commands::{CompleteWorkflowExecution, QueryResult, QuerySuccess},
+        workflow_completion::WorkflowActivationCompletion,
     },
     temporal::api::{
         failure::v1::Failure,
@@ -21,6 +25,7 @@ use temporal_sdk_core_protos::{
         },
     },
 };
+use temporal_sdk_core_test_utils::start_timer_cmd;
 
 #[rstest::rstest]
 #[case::with_history(true)]
@@ -37,6 +42,7 @@ async fn legacy_query(#[case] include_history: bool) {
             pr.query = Some(WorkflowQuery {
                 query_type: "query-type".to_string(),
                 query_args: Some(b"hi".into()),
+                header: None,
             });
             if !include_history {
                 pr.history = Some(History { events: vec![] });
@@ -45,7 +51,7 @@ async fn legacy_query(#[case] include_history: bool) {
         },
         hist_to_poll_resp(&t, wfid.to_owned(), 2.into(), TEST_Q.to_string()),
     ]);
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     mock_gateway
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
@@ -62,21 +68,17 @@ async fn legacy_query(#[case] include_history: bool) {
 
     let first_wft = || async {
         let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
-        core.complete_workflow_activation(WfActivationCompletion::from_cmd(
+        core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
             TEST_Q,
             task.run_id,
-            StartTimer {
-                seq: 1,
-                ..Default::default()
-            }
-            .into(),
+            start_timer_cmd(1, Duration::from_secs(1)),
         ))
         .await
         .unwrap();
     };
     let clear_eviction = || async {
         let t = core.poll_workflow_activation(TEST_Q).await.unwrap();
-        core.complete_workflow_activation(WfActivationCompletion::empty(TEST_Q, t.run_id))
+        core.complete_workflow_activation(WorkflowActivationCompletion::empty(TEST_Q, t.run_id))
             .await
             .unwrap();
     };
@@ -92,12 +94,12 @@ async fn legacy_query(#[case] include_history: bool) {
     // Poll again, and we end up getting a `query` field query response
     let query = assert_matches!(
         task.jobs.as_slice(),
-        [WfActivationJob {
-            variant: Some(wf_activation_job::Variant::QueryWorkflow(q)),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::QueryWorkflow(q)),
         }] => q
     );
     // Complete the query
-    core.complete_workflow_activation(WfActivationCompletion::from_cmd(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
         TEST_Q,
         task.run_id,
         QueryResult {
@@ -122,11 +124,11 @@ async fn legacy_query(#[case] include_history: bool) {
     let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
     assert_matches!(
         task.jobs.as_slice(),
-        [WfActivationJob {
-            variant: Some(wf_activation_job::Variant::FireTimer(_)),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::FireTimer(_)),
         }]
     );
-    core.complete_workflow_activation(WfActivationCompletion::from_cmds(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
         TEST_Q,
         task.run_id,
         vec![CompleteWorkflowExecution { result: None }.into()],
@@ -155,13 +157,14 @@ async fn new_queries(#[case] num_queries: usize) {
                     WorkflowQuery {
                         query_type: "query-type".to_string(),
                         query_args: Some(b"hi".into()),
+                        header: None,
                     },
                 );
             }
             pr
         },
     ]);
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     mock_gateway
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
@@ -172,14 +175,10 @@ async fn new_queries(#[case] num_queries: usize) {
     let core = mock_core(mock);
 
     let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
-    core.complete_workflow_activation(WfActivationCompletion::from_cmd(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
         TEST_Q,
         task.run_id,
-        StartTimer {
-            seq: 1,
-            ..Default::default()
-        }
-        .into(),
+        start_timer_cmd(1, Duration::from_secs(1)),
     ))
     .await
     .unwrap();
@@ -187,15 +186,15 @@ async fn new_queries(#[case] num_queries: usize) {
     let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
     assert_matches!(
         task.jobs[0],
-        WfActivationJob {
-            variant: Some(wf_activation_job::Variant::FireTimer(_)),
+        WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::FireTimer(_)),
         }
     );
     for i in 1..=num_queries {
         assert_matches!(
             task.jobs[i],
-            WfActivationJob {
-                variant: Some(wf_activation_job::Variant::QueryWorkflow(_)),
+            WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::QueryWorkflow(_)),
             }
         );
     }
@@ -214,7 +213,7 @@ async fn new_queries(#[case] num_queries: usize) {
         })
         .collect();
     qresults.push(CompleteWorkflowExecution { result: None }.into());
-    core.complete_workflow_activation(WfActivationCompletion::from_cmds(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
         TEST_Q,
         task.run_id,
         qresults,
@@ -235,13 +234,14 @@ async fn legacy_query_failure_on_wft_failure() {
             pr.query = Some(WorkflowQuery {
                 query_type: "query-type".to_string(),
                 query_args: Some(b"hi".into()),
+                header: None,
             });
             pr.history = Some(History { events: vec![] });
             pr
         },
         hist_to_poll_resp(&t, wfid.to_owned(), 2.into(), TEST_Q.to_string()),
     ]);
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     mock_gateway
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
@@ -255,14 +255,10 @@ async fn legacy_query_failure_on_wft_failure() {
     let core = mock_core(mock);
 
     let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
-    core.complete_workflow_activation(WfActivationCompletion::from_cmd(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
         TEST_Q,
         task.run_id,
-        StartTimer {
-            seq: 1,
-            ..Default::default()
-        }
-        .into(),
+        start_timer_cmd(1, Duration::from_secs(1)),
     ))
     .await
     .unwrap();
@@ -271,12 +267,12 @@ async fn legacy_query_failure_on_wft_failure() {
     // Poll again, and we end up getting a `query` field query response
     assert_matches!(
         task.jobs.as_slice(),
-        [WfActivationJob {
-            variant: Some(wf_activation_job::Variant::QueryWorkflow(q)),
+        [WorkflowActivationJob {
+            variant: Some(workflow_activation_job::Variant::QueryWorkflow(q)),
         }] => q
     );
     // Fail wft which should result in query being failed
-    core.complete_workflow_activation(WfActivationCompletion::fail(
+    core.complete_workflow_activation(WorkflowActivationCompletion::fail(
         TEST_Q,
         task.run_id,
         Failure {
@@ -290,10 +286,17 @@ async fn legacy_query_failure_on_wft_failure() {
     core.shutdown().await;
 }
 
+#[rstest::rstest]
 #[tokio::test]
-async fn legacy_query_with_full_history_after_complete() {
+async fn legacy_query_after_complete(#[values(false, true)] full_history: bool) {
     let wfid = "fake_wf_id";
-    let t = canned_histories::single_timer_wf_completes("1");
+    let t = if full_history {
+        canned_histories::single_timer_wf_completes("1")
+    } else {
+        let mut t = canned_histories::single_timer("1");
+        t.add_workflow_task_completed();
+        t
+    };
     let query_with_hist_task = {
         let mut pr = hist_to_poll_resp(
             &t,
@@ -304,6 +307,7 @@ async fn legacy_query_with_full_history_after_complete() {
         pr.query = Some(WorkflowQuery {
             query_type: "query-type".to_string(),
             query_args: Some(b"hi".into()),
+            header: None,
         });
         pr
     };
@@ -317,7 +321,7 @@ async fn legacy_query_with_full_history_after_complete() {
         query_with_hist_task.clone(),
         query_with_hist_task,
     ]);
-    let mut mock_gateway = MockServerGatewayApis::new();
+    let mut mock_gateway = mock_gateway();
     mock_gateway
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
@@ -331,19 +335,15 @@ async fn legacy_query_with_full_history_after_complete() {
     let core = mock_core(mock);
 
     let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
-    core.complete_workflow_activation(WfActivationCompletion::from_cmd(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
         TEST_Q,
         task.run_id,
-        StartTimer {
-            seq: 1,
-            ..Default::default()
-        }
-        .into(),
+        start_timer_cmd(1, Duration::from_secs(1)),
     ))
     .await
     .unwrap();
     let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
-    core.complete_workflow_activation(WfActivationCompletion::from_cmds(
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
         TEST_Q,
         task.run_id,
         vec![CompleteWorkflowExecution { result: None }.into()],
@@ -356,11 +356,11 @@ async fn legacy_query_with_full_history_after_complete() {
         let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
         let query = assert_matches!(
             task.jobs.as_slice(),
-            [WfActivationJob {
-                variant: Some(wf_activation_job::Variant::QueryWorkflow(q)),
+            [WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::QueryWorkflow(q)),
             }] => q
         );
-        core.complete_workflow_activation(WfActivationCompletion::from_cmd(
+        core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
             TEST_Q,
             task.run_id,
             QueryResult {