@temporalio/core-bridge 1.15.0 → 1.16.0

package/sdk-core/crates/sdk-core/tests/integ_tests/workflow_tests/patches.rs

@@ -7,59 +7,73 @@ use std::{
     },
     time::Duration,
 };
-use temporalio_client::WorkflowClientTrait;
-use temporalio_common::protos::{
-    DEFAULT_WORKFLOW_TYPE, TestHistoryBuilder, VERSION_SEARCH_ATTR_KEY,
-    constants::PATCH_MARKER_NAME,
-    coresdk::{
-        AsJsonPayloadExt, FromJsonPayloadExt,
-        common::decode_change_marker_details,
-        workflow_activation::{NotifyHasPatch, WorkflowActivationJob, workflow_activation_job},
-    },
-    temporal::api::{
-        command::v1::{
-            RecordMarkerCommandAttributes, ScheduleActivityTaskCommandAttributes,
-            UpsertWorkflowSearchAttributesCommandAttributes, command::Attributes,
+use temporalio_client::{WorkflowSignalOptions, WorkflowStartOptions};
+use temporalio_common::{
+    data_converters::RawValue,
+    protos::{
+        DEFAULT_WORKFLOW_TYPE, TestHistoryBuilder, VERSION_SEARCH_ATTR_KEY,
+        constants::PATCH_MARKER_NAME,
+        coresdk::{
+            AsJsonPayloadExt, FromJsonPayloadExt,
+            common::decode_change_marker_details,
+            workflow_activation::{NotifyHasPatch, WorkflowActivationJob, workflow_activation_job},
         },
-        common::v1::ActivityType,
-        enums::v1::{CommandType, EventType},
-        history::v1::{
-            ActivityTaskCompletedEventAttributes, ActivityTaskScheduledEventAttributes,
-            ActivityTaskStartedEventAttributes, TimerFiredEventAttributes,
+        temporal::api::{
+            command::v1::{
+                RecordMarkerCommandAttributes, ScheduleActivityTaskCommandAttributes,
+                UpsertWorkflowSearchAttributesCommandAttributes, command::Attributes,
+            },
+            common::v1::ActivityType,
+            enums::v1::{CommandType, EventType},
+            history::v1::{
+                ActivityTaskCompletedEventAttributes, ActivityTaskScheduledEventAttributes,
+                ActivityTaskStartedEventAttributes, TimerFiredEventAttributes,
+            },
         },
     },
 };
 
 use temporalio_common::worker::WorkerTaskTypes;
-use temporalio_sdk::{ActivityOptions, WfContext, WorkflowResult};
+use temporalio_macros::{activities, workflow, workflow_methods};
+use temporalio_sdk::{
+    ActivityOptions, SyncWorkflowContext, WorkflowContext, WorkflowResult,
+    activities::{ActivityContext, ActivityError},
+};
 use temporalio_sdk_core::test_help::{CoreInternalFlags, MockPollCfg, ResponseType};
 use tokio::{join, sync::Notify};
-use tokio_stream::StreamExt;
 
 const MY_PATCH_ID: &str = "integ_test_change_name";
 
-pub(crate) async fn changes_wf(ctx: WfContext) -> WorkflowResult<()> {
-    if ctx.patched(MY_PATCH_ID) {
-        ctx.timer(Duration::from_millis(100)).await;
-    } else {
-        ctx.timer(Duration::from_millis(200)).await;
-    }
-    ctx.timer(Duration::from_millis(200)).await;
-    if ctx.patched(MY_PATCH_ID) {
-        ctx.timer(Duration::from_millis(100)).await;
-    } else {
+#[workflow]
+#[derive(Default)]
+pub(crate) struct ChangesWf;
+
+#[workflow_methods]
+impl ChangesWf {
+    #[run(name = "writes_change_markers")]
+    pub(crate) async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        if ctx.patched(MY_PATCH_ID) {
+            ctx.timer(Duration::from_millis(100)).await;
+        } else {
+            ctx.timer(Duration::from_millis(200)).await;
+        }
         ctx.timer(Duration::from_millis(200)).await;
+        if ctx.patched(MY_PATCH_ID) {
+            ctx.timer(Duration::from_millis(100)).await;
+        } else {
+            ctx.timer(Duration::from_millis(200)).await;
+        }
+        Ok(())
     }
-    Ok(().into())
 }
 
 #[tokio::test]
 async fn writes_change_markers() {
     let wf_name = "writes_change_markers";
     let mut starter = CoreWfStarter::new(wf_name);
-    starter.worker_config.task_types = WorkerTaskTypes::workflow_only();
+    starter.sdk_config.task_types = WorkerTaskTypes::workflow_only();
     let mut worker = starter.worker().await;
-    worker.register_wf(wf_name.to_owned(), changes_wf);
+    worker.register_workflow::<ChangesWf>();
 
     starter.start_with_worker(wf_name, &mut worker).await;
     worker.run_until_done().await.unwrap();
@@ -67,57 +81,77 @@ async fn writes_change_markers() {
 
 /// This one simulates a run as if the worker had the "old" code, then it fails at the end as
 /// a cheapo way of being re-run, at which point it runs with change checks and the "new" code.
-static DID_DIE: AtomicBool = AtomicBool::new(false);
+#[workflow]
+pub(crate) struct NoChangeThenChangeWf {
+    did_die: Arc<AtomicBool>,
+}
 
-pub(crate) async fn no_change_then_change_wf(ctx: WfContext) -> WorkflowResult<()> {
-    if DID_DIE.load(Ordering::Acquire) {
-        assert!(!ctx.patched(MY_PATCH_ID));
-    }
-    ctx.timer(Duration::from_millis(200)).await;
-    ctx.timer(Duration::from_millis(200)).await;
-    if DID_DIE.load(Ordering::Acquire) {
-        assert!(!ctx.patched(MY_PATCH_ID));
-    }
-    ctx.timer(Duration::from_millis(200)).await;
+#[workflow_methods(factory_only)]
+impl NoChangeThenChangeWf {
+    #[run(name = "can_add_change_markers")]
+    pub(crate) async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        if ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            assert!(!ctx.patched(MY_PATCH_ID));
+        }
+        ctx.timer(Duration::from_millis(200)).await;
+        ctx.timer(Duration::from_millis(200)).await;
+        if ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            assert!(!ctx.patched(MY_PATCH_ID));
+        }
+        ctx.timer(Duration::from_millis(200)).await;
 
-    if !DID_DIE.load(Ordering::Acquire) {
-        DID_DIE.store(true, Ordering::Release);
-        ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
+        if !ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            ctx.state(|wf| wf.did_die.store(true, Ordering::Release));
+            ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
+        }
+        Ok(())
     }
-    Ok(().into())
 }
 
 #[tokio::test]
 async fn can_add_change_markers() {
     let wf_name = "can_add_change_markers";
     let mut starter = CoreWfStarter::new(wf_name);
-    starter.worker_config.task_types = WorkerTaskTypes::workflow_only();
+    starter.sdk_config.task_types = WorkerTaskTypes::workflow_only();
     let mut worker = starter.worker().await;
-    worker.register_wf(wf_name.to_owned(), no_change_then_change_wf);
+    let did_die = Arc::new(AtomicBool::new(false));
+    worker.register_workflow_with_factory(move || NoChangeThenChangeWf {
+        did_die: did_die.clone(),
+    });
 
     starter.start_with_worker(wf_name, &mut worker).await;
     worker.run_until_done().await.unwrap();
 }
 
-static DID_DIE_2: AtomicBool = AtomicBool::new(false);
+#[workflow]
+pub(crate) struct ReplayWithChangeMarkerWf {
+    did_die: Arc<AtomicBool>,
+}
 
-pub(crate) async fn replay_with_change_marker_wf(ctx: WfContext) -> WorkflowResult<()> {
-    assert!(ctx.patched(MY_PATCH_ID));
-    ctx.timer(Duration::from_millis(200)).await;
-    if !DID_DIE_2.load(Ordering::Acquire) {
-        DID_DIE_2.store(true, Ordering::Release);
-        ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
+#[workflow_methods(factory_only)]
+impl ReplayWithChangeMarkerWf {
+    #[run(name = "replaying_with_patch_marker")]
+    pub(crate) async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        assert!(ctx.patched(MY_PATCH_ID));
+        ctx.timer(Duration::from_millis(200)).await;
+        if !ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            ctx.state(|wf| wf.did_die.store(true, Ordering::Release));
+            ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
+        }
+        Ok(())
     }
-    Ok(().into())
 }
 
 #[tokio::test]
 async fn replaying_with_patch_marker() {
     let wf_name = "replaying_with_patch_marker";
     let mut starter = CoreWfStarter::new(wf_name);
-    starter.worker_config.task_types = WorkerTaskTypes::workflow_only();
+    starter.sdk_config.task_types = WorkerTaskTypes::workflow_only();
     let mut worker = starter.worker().await;
-    worker.register_wf(wf_name.to_owned(), replay_with_change_marker_wf);
+    let did_die = Arc::new(AtomicBool::new(false));
+    worker.register_workflow_with_factory(move || ReplayWithChangeMarkerWf {
+        did_die: did_die.clone(),
+    });
 
     starter.start_with_worker(wf_name, &mut worker).await;
     worker.run_until_done().await.unwrap();
@@ -126,98 +160,148 @@ async fn replaying_with_patch_marker() {
 /// Test that the internal patching mechanism works on the second workflow task when replaying.
 /// Used as regression test for a bug that detected that we did not look ahead far enough to find
 /// the next workflow task completion, which the flags are attached to.
+#[workflow]
+struct TimerPatchedTimerWf {
+    fail_once: Arc<AtomicBool>,
+}
+
+#[workflow_methods(factory_only)]
+impl TimerPatchedTimerWf {
+    #[run(name = "timer_patched_timer")]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        ctx.timer(Duration::from_millis(1)).await;
+        if ctx.state(|wf| wf.fail_once.load(Ordering::Acquire)) {
+            ctx.state(|wf| wf.fail_once.store(false, Ordering::Release));
+            panic!("Enchi is hungry!");
+        }
+        assert!(ctx.patched(MY_PATCH_ID));
+        ctx.timer(Duration::from_millis(1)).await;
+        Ok(())
+    }
+}
+
 #[tokio::test]
 async fn patched_on_second_workflow_task_is_deterministic() {
     let wf_name = "timer_patched_timer";
     let mut starter = CoreWfStarter::new(wf_name);
     // Disable caching to force replay from beginning
-    starter.worker_config.max_cached_workflows = 0_usize;
-    starter.worker_config.task_types = WorkerTaskTypes::workflow_only();
+    starter.sdk_config.max_cached_workflows = 0_usize;
+    starter.sdk_config.task_types = WorkerTaskTypes::workflow_only();
     let mut worker = starter.worker().await;
-    // Include a task failure as well to make sure that works
-    static FAIL_ONCE: AtomicBool = AtomicBool::new(true);
-    worker.register_wf(wf_name.to_owned(), |ctx: WfContext| async move {
-        ctx.timer(Duration::from_millis(1)).await;
-        if FAIL_ONCE.load(Ordering::Acquire) {
-            FAIL_ONCE.store(false, Ordering::Release);
-            panic!("Enchi is hungry!");
-        }
-        assert!(ctx.patched(MY_PATCH_ID));
-        ctx.timer(Duration::from_millis(1)).await;
-        Ok(().into())
+    let fail_once = Arc::new(AtomicBool::new(true));
+    worker.register_workflow_with_factory(move || TimerPatchedTimerWf {
+        fail_once: fail_once.clone(),
     });
 
     starter.start_with_worker(wf_name, &mut worker).await;
     worker.run_until_done().await.unwrap();
 }
 
+#[workflow]
+struct RemoveDeprecatedPatchNearOtherPatchWf {
+    did_die: Arc<AtomicBool>,
+}
+
+#[workflow_methods(factory_only)]
+impl RemoveDeprecatedPatchNearOtherPatchWf {
+    #[run(name = "can_add_change_markers")]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        ctx.timer(Duration::from_millis(200)).await;
+        if !ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            assert!(ctx.deprecate_patch("getting-deprecated"));
+            assert!(ctx.patched("staying"));
+        } else {
+            assert!(ctx.patched("staying"));
+        }
+        ctx.timer(Duration::from_millis(200)).await;
+
+        if !ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            ctx.state(|wf| wf.did_die.store(true, Ordering::Release));
+            ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
+        }
+        Ok(())
+    }
+}
+
 #[tokio::test]
 async fn can_remove_deprecated_patch_near_other_patch() {
     let wf_name = "can_add_change_markers";
     let mut starter = CoreWfStarter::new(wf_name);
-    starter.worker_config.task_types = WorkerTaskTypes::workflow_only();
+    starter.sdk_config.task_types = WorkerTaskTypes::workflow_only();
     let mut worker = starter.worker().await;
     let did_die = Arc::new(AtomicBool::new(false));
-    worker.register_wf(wf_name.to_owned(), move |ctx: WfContext| {
-        let did_die = did_die.clone();
-        async move {
-            ctx.timer(Duration::from_millis(200)).await;
-            if !did_die.load(Ordering::Acquire) {
-                assert!(ctx.deprecate_patch("getting-deprecated"));
-                assert!(ctx.patched("staying"));
-            } else {
-                assert!(ctx.patched("staying"));
-            }
-            ctx.timer(Duration::from_millis(200)).await;
-
-            if !did_die.load(Ordering::Acquire) {
-                did_die.store(true, Ordering::Release);
-                ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
-            }
-            Ok(().into())
-        }
+    worker.register_workflow_with_factory(move || RemoveDeprecatedPatchNearOtherPatchWf {
+        did_die: did_die.clone(),
     });
 
     starter.start_with_worker(wf_name, &mut worker).await;
     worker.run_until_done().await.unwrap();
 }
 
+#[workflow]
+struct DeprecatedPatchRemovalWf {
+    did_die: Arc<AtomicBool>,
+    notify: Arc<Notify>,
+    signal_received: bool,
+}
+
+#[workflow_methods(factory_only)]
+impl DeprecatedPatchRemovalWf {
+    #[run(name = "deprecated_patch_removal")]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        if !ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            assert!(ctx.deprecate_patch("getting-deprecated"));
+        }
+        ctx.state(|wf| wf.notify.notify_one());
+        ctx.wait_condition(|s| s.signal_received).await;
+
+        ctx.timer(Duration::from_millis(1)).await;
+
+        if !ctx.state(|wf| wf.did_die.load(Ordering::Acquire)) {
+            ctx.state(|wf| wf.did_die.store(true, Ordering::Release));
+            ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
+        }
+        Ok(())
+    }
+
+    #[signal]
+    fn handle_sig(&mut self, _ctx: &mut SyncWorkflowContext<Self>) {
+        self.signal_received = true;
+    }
+}
+
 #[tokio::test]
 async fn deprecated_patch_removal() {
     let wf_name = "deprecated_patch_removal";
     let mut starter = CoreWfStarter::new(wf_name);
-    starter.worker_config.task_types = WorkerTaskTypes::workflow_only();
+    starter.sdk_config.task_types = WorkerTaskTypes::workflow_only();
     let mut worker = starter.worker().await;
-    let client = starter.get_client().await;
     let wf_id = starter.get_task_queue().to_string();
     let did_die = Arc::new(AtomicBool::new(false));
     let send_sig = Arc::new(Notify::new());
-    let send_sig_c = send_sig.clone();
-    worker.register_wf(wf_name, move |ctx: WfContext| {
-        let did_die = did_die.clone();
-        let send_sig_c = send_sig_c.clone();
-        async move {
-            if !did_die.load(Ordering::Acquire) {
-                assert!(ctx.deprecate_patch("getting-deprecated"));
-            }
-            send_sig_c.notify_one();
-            ctx.make_signal_channel("sig").next().await;
-
-            ctx.timer(Duration::from_millis(1)).await;
-
-            if !did_die.load(Ordering::Acquire) {
-                did_die.store(true, Ordering::Release);
-                ctx.force_task_fail(anyhow::anyhow!("i'm ded"));
-            }
-            Ok(().into())
-        }
+    let send_sig_clone = send_sig.clone();
+    worker.register_workflow_with_factory(move || DeprecatedPatchRemovalWf {
+        did_die: did_die.clone(),
+        notify: send_sig_clone.clone(),
+        signal_received: false,
     });
 
-    starter.start_with_worker(wf_name, &mut worker).await;
+    let handle = worker
+        .submit_workflow(
+            DeprecatedPatchRemovalWf::run,
+            (),
+            WorkflowStartOptions::new(wf_id.clone(), wf_id).build(),
+        )
+        .await
+        .unwrap();
     let sig_fut = async {
         send_sig.notified().await;
-        client
-            .signal_workflow_execution(wf_id, "".to_string(), "sig".to_string(), None, None)
+        handle
+            .signal(
+                DeprecatedPatchRemovalWf::handle_sig,
+                (),
+                WorkflowSignalOptions::default(),
+            )
             .await
             .unwrap()
     };
@@ -312,47 +396,81 @@ fn patch_marker_single_activity(
     t
 }
 
-async fn v1(ctx: &mut WfContext) {
-    ctx.activity(ActivityOptions {
-        activity_id: Some("no_change".to_owned()),
-        ..Default::default()
-    })
-    .await;
+struct FakeAct;
+#[activities]
+impl FakeAct {
+    #[activity(name = "")]
+    fn nameless(_: ActivityContext) -> Result<RawValue, ActivityError> {
+        unimplemented!()
+    }
 }
 
-async fn v2(ctx: &mut WfContext) -> bool {
-    if ctx.patched(MY_PATCH_ID) {
-        ctx.activity(ActivityOptions {
-            activity_id: Some("had_change".to_owned()),
-            ..Default::default()
-        })
+async fn v1(ctx: &mut WorkflowContext<PatchWf>) {
+    let _ = ctx
+        .start_activity(
+            FakeAct::nameless,
+            (),
+            ActivityOptions {
+                activity_id: Some("no_change".to_owned()),
+                ..Default::default()
+            },
+        )
         .await;
+}
+
+async fn v2(ctx: &mut WorkflowContext<PatchWf>) -> bool {
+    if ctx.patched(MY_PATCH_ID) {
+        let _ = ctx
+            .start_activity(
+                FakeAct::nameless,
+                (),
+                ActivityOptions {
+                    activity_id: Some("had_change".to_owned()),
+                    ..Default::default()
+                },
+            )
+            .await;
         true
     } else {
-        ctx.activity(ActivityOptions {
-            activity_id: Some("no_change".to_owned()),
-            ..Default::default()
-        })
-        .await;
+        let _ = ctx
+            .start_activity(
+                FakeAct::nameless,
+                (),
+                ActivityOptions {
+                    activity_id: Some("no_change".to_owned()),
+                    ..Default::default()
+                },
+            )
+            .await;
         false
     }
 }
 
-async fn v3(ctx: &mut WfContext) {
+async fn v3(ctx: &mut WorkflowContext<PatchWf>) {
     ctx.deprecate_patch(MY_PATCH_ID);
-    ctx.activity(ActivityOptions {
-        activity_id: Some("had_change".to_owned()),
-        ..Default::default()
-    })
-    .await;
+    let _ = ctx
+        .start_activity(
+            FakeAct::nameless,
+            (),
+            ActivityOptions {
+                activity_id: Some("had_change".to_owned()),
+                ..Default::default()
+            },
+        )
+        .await;
 }
 
-async fn v4(ctx: &mut WfContext) {
-    ctx.activity(ActivityOptions {
-        activity_id: Some("had_change".to_owned()),
-        ..Default::default()
-    })
-    .await;
+async fn v4(ctx: &mut WorkflowContext<PatchWf>) {
+    let _ = ctx
+        .start_activity(
+            FakeAct::nameless,
+            (),
+            ActivityOptions {
+                activity_id: Some("had_change".to_owned()),
+                ..Default::default()
+            },
+        )
+        .await;
 }
 
 fn patch_setup(replaying: bool, marker_type: MarkerType, workflow_version: usize) -> MockPollCfg {
@@ -364,27 +482,32 @@ fn patch_setup(replaying: bool, marker_type: MarkerType, workflow_version: usize
     }
 }
 
-macro_rules! patch_wf {
-    ($workflow_version:ident) => {
-        move |mut ctx: WfContext| async move {
-            match $workflow_version {
-                1 => {
-                    v1(&mut ctx).await;
-                }
-                2 => {
-                    v2(&mut ctx).await;
-                }
-                3 => {
-                    v3(&mut ctx).await;
-                }
-                4 => {
-                    v4(&mut ctx).await;
-                }
-                _ => panic!("Invalid workflow version for test setup"),
+#[workflow]
+struct PatchWf {
+    version: usize,
+}
+
+#[workflow_methods(factory_only)]
+impl PatchWf {
+    #[run(name = DEFAULT_WORKFLOW_TYPE)]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        match ctx.state(|wf| wf.version) {
+            1 => {
+                v1(ctx).await;
+            }
+            2 => {
+                v2(ctx).await;
+            }
+            3 => {
+                v3(ctx).await;
             }
-            Ok(().into())
+            4 => {
+                v4(ctx).await;
+            }
+            _ => panic!("Invalid workflow version for test setup"),
         }
-    };
+        Ok(())
+    }
 }
 
 #[rstest]
@@ -436,7 +559,9 @@ async fn v1_and_v4_changes(
 
     let mut worker = build_fake_sdk(mock_cfg);
     worker.set_worker_interceptor(aai);
-    worker.register_wf(DEFAULT_WORKFLOW_TYPE, patch_wf!(wf_version));
+    worker.register_workflow_with_factory(move || PatchWf {
+        version: wf_version,
+    });
     worker.run().await.unwrap();
 }
 
@@ -541,10 +666,39 @@ async fn v2_and_v3_changes(
 
     let mut worker = build_fake_sdk(mock_cfg);
     worker.set_worker_interceptor(aai);
-    worker.register_wf(DEFAULT_WORKFLOW_TYPE, patch_wf!(wf_version));
+    worker.register_workflow_with_factory(move || PatchWf {
+        version: wf_version,
+    });
     worker.run().await.unwrap();
 }
 
+#[workflow]
+#[derive(Default)]
+struct SameChangeMultipleSpotsWf;
+
+#[workflow_methods]
+impl SameChangeMultipleSpotsWf {
+    #[run(name = DEFAULT_WORKFLOW_TYPE)]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        if ctx.patched(MY_PATCH_ID) {
+            let _ = ctx
+                .start_activity(FakeAct::nameless, (), ActivityOptions::default())
+                .await;
+        } else {
+            ctx.timer(ONE_SECOND).await;
+        }
+        ctx.timer(ONE_SECOND).await;
+        if ctx.patched(MY_PATCH_ID) {
+            let _ = ctx
+                .start_activity(FakeAct::nameless, (), ActivityOptions::default())
+                .await;
+        } else {
+            ctx.timer(ONE_SECOND).await;
+        }
+        Ok(())
+    }
+}
+
 #[rstest]
 #[case::has_change_replay(true, true)]
 #[case::no_change_replay(false, true)]
@@ -634,26 +788,30 @@ async fn same_change_multiple_spots(#[case] have_marker_in_hist: bool, #[case] r
         MockPollCfg::from_hist_builder(t)
     };
 
-    // Errors would appear as nondeterminism problems, so just run it.
     let mut worker = build_fake_sdk(mock_cfg);
-    worker.register_wf(DEFAULT_WORKFLOW_TYPE, move |ctx: WfContext| async move {
-        if ctx.patched(MY_PATCH_ID) {
-            ctx.activity(ActivityOptions::default()).await;
-        } else {
-            ctx.timer(ONE_SECOND).await;
-        }
-        ctx.timer(ONE_SECOND).await;
-        if ctx.patched(MY_PATCH_ID) {
-            ctx.activity(ActivityOptions::default()).await;
-        } else {
-            ctx.timer(ONE_SECOND).await;
-        }
-        Ok(().into())
-    });
+    worker.register_workflow::<SameChangeMultipleSpotsWf>();
     worker.run().await.unwrap();
 }
 
 const SIZE_OVERFLOW_PATCH_AMOUNT: usize = 180;
+
+#[workflow]
+struct ManyPatchesWf {
+    num_patches: usize,
+}
+
+#[workflow_methods(factory_only)]
+impl ManyPatchesWf {
+    #[run(name = DEFAULT_WORKFLOW_TYPE)]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        for i in 1..=ctx.state(|wf| wf.num_patches) {
+            let _dontcare = ctx.patched(&format!("patch-{i}"));
+            ctx.timer(ONE_SECOND).await;
+        }
+        Ok(())
+    }
+}
+
 #[rstest]
 #[case::happy_path(50)]
 // We start exceeding the 2k size limit at 180 patches with this format
@@ -713,12 +871,6 @@ async fn many_patches_combine_in_search_attrib_update(#[case] num_patches: usize
     });
 
     let mut worker = build_fake_sdk(mock_cfg);
-    worker.register_wf(DEFAULT_WORKFLOW_TYPE, move |ctx: WfContext| async move {
-        for i in 1..=num_patches {
-            let _dontcare = ctx.patched(&format!("patch-{i}"));
-            ctx.timer(ONE_SECOND).await;
-        }
-        Ok(().into())
-    });
+    worker.register_workflow_with_factory(move || ManyPatchesWf { num_patches });
     worker.run().await.unwrap();
 }