@temporalio/core-bridge 1.13.1 → 1.14.0

package/sdk-core/{core/src/test_help/mod.rs → crates/sdk-core/src/test_help/integ_helpers.rs}

@@ -1,26 +1,31 @@
-pub(crate) use temporal_sdk_core_test_utils::canned_histories;
+//! Integration test helpers - available with test-utilities feature or in test mode
+
+#![allow(missing_docs)]
+
+pub use crate::{
+    internal_flags::CoreInternalFlags,
+    worker::{LEGACY_QUERY_ID, client::mocks::mock_worker_client},
+};
 
 use crate::{
     TaskToken, Worker, WorkerConfig, WorkerConfigBuilder,
-    pollers::{BoxedPoller, MockManualPoller, MockPoller},
+    pollers::{BoxedPoller, MockManualPoller},
     protosext::ValidPollWFTQResponse,
     replay::TestHistoryBuilder,
     sticky_q_name_for_worker,
     worker::{
         TaskPollers,
-        client::{
-            LegacyQueryResult, MockWorkerClient, WorkerClient, WorkflowTaskCompletion,
-            mocks::mock_worker_client,
-        },
+        client::{LegacyQueryResult, MockWorkerClient, WorkerClient, WorkflowTaskCompletion},
     },
 };
+use assert_matches::assert_matches;
 use async_trait::async_trait;
 use bimap::BiMap;
 use futures_util::{FutureExt, Stream, StreamExt, future::BoxFuture, stream, stream::BoxStream};
 use mockall::TimesRange;
 use parking_lot::RwLock;
 use std::{
-    collections::{BTreeMap, HashMap, HashSet, VecDeque},
+    collections::{BTreeMap, HashMap, VecDeque},
     fmt::Debug,
     ops::{Deref, DerefMut},
     pin::Pin,
@@ -31,46 +36,76 @@ use std::{
     task::{Context, Poll},
     time::Duration,
 };
-use temporal_sdk::interceptors::FailOnNondeterminismInterceptor;
-use temporal_sdk_core_api::{
+use temporalio_common::{
     Worker as WorkerTrait,
     errors::PollError,
-    worker::{PollerBehavior, WorkerVersioningStrategy},
-};
-use temporal_sdk_core_protos::{
-    coresdk::{
-        workflow_activation::{WorkflowActivation, workflow_activation_job},
-        workflow_commands::workflow_command,
-        workflow_completion::{self, WorkflowActivationCompletion, workflow_activation_completion},
-    },
-    temporal::api::{
-        common::v1::WorkflowExecution,
-        enums::v1::WorkflowTaskFailedCause,
-        failure::v1::Failure,
-        protocol,
-        protocol::v1::message,
-        update,
-        workflowservice::v1::{
-            PollActivityTaskQueueResponse, PollNexusTaskQueueResponse,
-            PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse,
+    protos::{
+        coresdk::{
+            workflow_activation::{WorkflowActivationJob, workflow_activation_job},
+            workflow_commands::{CompleteWorkflowExecution, StartTimer},
+            workflow_completion::WorkflowActivationCompletion,
         },
+        temporal::api::{
+            common::v1::WorkflowExecution,
+            enums::v1::WorkflowTaskFailedCause,
+            failure::v1::Failure,
+            protocol::{self, v1::message},
+            update,
+            workflowservice::v1::{
+                PollActivityTaskQueueResponse, PollNexusTaskQueueResponse,
+                PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse,
+            },
+        },
+        utilities::pack_any,
     },
-    utilities::pack_any,
+    worker::{PollerBehavior, WorkerTaskTypes, WorkerVersioningStrategy},
 };
-use temporal_sdk_core_test_utils::{NAMESPACE, TestWorker};
 use tokio::sync::{Notify, mpsc::unbounded_channel};
 use tokio_stream::wrappers::UnboundedReceiverStream;
+use uuid::Uuid;
+
+/// Default namespace for testing
+pub const NAMESPACE: &str = "default";
 
-pub(crate) const TEST_Q: &str = "q";
+/// Initiate shutdown, drain the pollers (handling evictions), and wait for shutdown to complete.
+pub async fn drain_pollers_and_shutdown(worker: &dyn WorkerTrait) {
+    worker.initiate_shutdown();
+    tokio::join!(
+        async {
+            assert_matches!(
+                worker.poll_activity_task().await.unwrap_err(),
+                PollError::ShutDown
+            );
+        },
+        async {
+            loop {
+                match worker.poll_workflow_activation().await {
+                    Err(PollError::ShutDown) => break,
+                    Ok(a) if a.is_only_eviction() => {
+                        worker
+                            .complete_workflow_activation(WorkflowActivationCompletion::empty(
+                                a.run_id,
+                            ))
+                            .await
+                            .unwrap();
+                    }
+                    o => panic!("Unexpected activation while draining: {o:?}"),
+                }
+            }
+        }
+    );
+    worker.shutdown().await;
+}
 
-pub(crate) fn test_worker_cfg() -> WorkerConfigBuilder {
+pub fn test_worker_cfg() -> WorkerConfigBuilder {
     let mut wcb = WorkerConfigBuilder::default();
     wcb.namespace(NAMESPACE)
-        .task_queue(TEST_Q)
+        .task_queue(Uuid::new_v4().to_string())
         .versioning_strategy(WorkerVersioningStrategy::None {
             build_id: "test_bin_id".to_string(),
         })
         .ignore_evicts_on_shutdown(true)
+        .task_types(WorkerTaskTypes::all())
         // Serial polling since it makes mocking much easier.
         .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize));
     wcb
@@ -79,7 +114,7 @@ pub(crate) fn test_worker_cfg() -> WorkerConfigBuilder {
 /// When constructing responses for mocks, indicates how a given response should be built
 #[derive(derive_more::From)]
 #[allow(clippy::large_enum_variant)] // Test only code, whatever.
-pub(crate) enum ResponseType {
+pub enum ResponseType {
     ToTaskNum(usize),
     /// Returns just the history from the WFT completed of the provided task number - 1, through to
     /// the next WFT started. Simulating the incremental history for just the provided task number
@@ -94,7 +129,7 @@ pub(crate) enum ResponseType {
 }
 
 #[derive(Eq, PartialEq, Hash)]
-pub(crate) enum HashableResponseType {
+pub enum HashableResponseType {
     ToTaskNum(usize),
     OneTask(usize),
     UntilResolved(usize),
@@ -130,7 +165,7 @@ impl From<&usize> for ResponseType {
 /// `response_batches` is used to control the fake [PollWorkflowTaskQueueResponse]s returned. For
 /// each number in the input list, a fake response will be prepared which includes history up to the
 /// workflow task with that number, as in [TestHistoryBuilder::get_history_info].
-pub(crate) fn build_fake_worker(
+pub fn build_fake_worker(
     wf_id: &str,
     t: TestHistoryBuilder,
     response_batches: impl IntoIterator<Item = impl Into<ResponseType>>,
@@ -148,24 +183,17 @@ pub(crate) fn build_fake_worker(
     mock_worker(mock_holder)
 }
 
-pub(crate) fn build_fake_sdk(mock_cfg: MockPollCfg) -> temporal_sdk::Worker {
-    let mut mock = build_mock_pollers(mock_cfg);
-    mock.worker_cfg(|c| {
-        c.max_cached_workflows = 1;
-        c.ignore_evicts_on_shutdown = false;
-    });
-    let core = mock_worker(mock);
-    let mut worker = temporal_sdk::Worker::new_from_core(Arc::new(core), "replay_q".to_string());
-    worker.set_worker_interceptor(FailOnNondeterminismInterceptor {});
-    worker
-}
-
-pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
-    let sticky_q = sticky_q_name_for_worker("unit-test", &mocks.inputs.config);
-    let act_poller = if mocks.inputs.config.no_remote_activities {
+pub fn mock_worker(mocks: MocksHolder) -> Worker {
+    let sticky_q = sticky_q_name_for_worker("unit-test", mocks.inputs.config.max_cached_workflows);
+    let act_poller = if mocks.inputs.config.task_types.enable_remote_activities {
+        mocks.inputs.act_poller
+    } else {
         None
+    };
+    let nexus_poller = if mocks.inputs.config.task_types.enable_nexus {
+        mocks.inputs.nexus_poller
     } else {
-        mocks.inputs.act_poller
+        None
     };
     Worker::new_with_pollers(
         mocks.inputs.config,
@@ -174,45 +202,30 @@ pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
         TaskPollers::Mocked {
             wft_stream: mocks.inputs.wft_stream,
             act_poller,
-            nexus_poller: mocks
-                .inputs
-                .nexus_poller
-                .unwrap_or_else(|| mock_poller_from_resps([])),
+            nexus_poller,
         },
         None,
         None,
+        false,
     )
+    .unwrap()
 }
 
-pub(crate) fn mock_sdk(poll_cfg: MockPollCfg) -> TestWorker {
-    mock_sdk_cfg(poll_cfg, |_| {})
-}
-pub(crate) fn mock_sdk_cfg(
-    mut poll_cfg: MockPollCfg,
-    mutator: impl FnOnce(&mut WorkerConfig),
-) -> TestWorker {
-    poll_cfg.using_rust_sdk = true;
-    let mut mock = build_mock_pollers(poll_cfg);
-    mock.worker_cfg(mutator);
-    let core = mock_worker(mock);
-    TestWorker::new(Arc::new(core), TEST_Q.to_string())
-}
-
-pub(crate) struct FakeWfResponses {
-    pub(crate) wf_id: String,
-    pub(crate) hist: TestHistoryBuilder,
-    pub(crate) response_batches: Vec<ResponseType>,
+pub struct FakeWfResponses {
+    pub wf_id: String,
+    pub hist: TestHistoryBuilder,
+    pub response_batches: Vec<ResponseType>,
 }
 
 // TODO: Should be all-internal to this module
-pub(crate) struct MocksHolder {
+pub struct MocksHolder {
     client: Arc<dyn WorkerClient>,
     inputs: MockWorkerInputs,
-    pub(crate) outstanding_task_map: Option<OutstandingWFTMap>,
+    pub outstanding_task_map: Option<OutstandingWFTMap>,
 }
 
 impl MocksHolder {
-    pub(crate) fn worker_cfg(&mut self, mutator: impl FnOnce(&mut WorkerConfig)) {
+    pub fn worker_cfg(&mut self, mutator: impl FnOnce(&mut WorkerConfig)) {
         mutator(&mut self.inputs.config);
     }
 
@@ -220,15 +233,38 @@ impl MocksHolder {
         self.inputs.act_poller = Some(poller);
     }
 
+    /// Helper to create and set an activity poller from responses
+    #[cfg(test)]
+    pub(crate) fn set_act_poller_from_resps<ACT>(&mut self, act_tasks: ACT)
+    where
+        ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
+        <ACT as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let act_poller = mock_poller_from_resps(act_tasks);
+        self.set_act_poller(act_poller);
+    }
+
+    /// Helper to create and set a nexus poller from responses
+    #[cfg(test)]
+    pub(crate) fn set_nexus_poller_from_resps<NEX>(&mut self, nexus_tasks: NEX)
+    where
+        NEX: IntoIterator<Item = QueueResponse<PollNexusTaskQueueResponse>>,
+        <NEX as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let nexus_poller = mock_poller_from_resps(nexus_tasks);
+        self.inputs.nexus_poller = Some(nexus_poller);
+    }
+
     /// Can be used for tests that need to avoid auto-shutdown due to running out of mock responses
-    pub(crate) fn make_wft_stream_interminable(&mut self) {
-        let old_stream = std::mem::replace(&mut self.inputs.wft_stream, stream::pending().boxed());
-        self.inputs.wft_stream = old_stream.chain(stream::pending()).boxed();
+    pub fn make_wft_stream_interminable(&mut self) {
+        if let Some(old_stream) = self.inputs.wft_stream.take() {
+            self.inputs.wft_stream = Some(old_stream.chain(stream::pending()).boxed());
+        }
     }
 }
 
-pub(crate) struct MockWorkerInputs {
-    pub(crate) wft_stream: BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>,
+pub struct MockWorkerInputs {
+    pub(crate) wft_stream: Option<BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>>,
     pub(crate) act_poller: Option<BoxedPoller<PollActivityTaskQueueResponse>>,
     pub(crate) nexus_poller: Option<BoxedPoller<PollNexusTaskQueueResponse>>,
     pub(crate) config: WorkerConfig,
@@ -245,7 +281,7 @@ impl MockWorkerInputs {
         wft_stream: BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>,
     ) -> Self {
         Self {
-            wft_stream,
+            wft_stream: Some(wft_stream),
             act_poller: None,
             nexus_poller: None,
             config: test_worker_cfg().build().unwrap(),
@@ -254,7 +290,7 @@ impl MockWorkerInputs {
 }
 
 impl MocksHolder {
-    pub(crate) fn from_mock_worker(
+    pub fn from_mock_worker(
         client: impl WorkerClient + 'static,
         mock_worker: MockWorkerInputs,
     ) -> Self {
@@ -265,8 +301,8 @@ impl MocksHolder {
         }
     }
 
-    /// Uses the provided list of tasks to create a mock poller for the `TEST_Q`
-    pub(crate) fn from_client_with_activities<ACT>(
+    /// Uses the provided list of tasks to create a mock poller with a randomly generated task queue
+    pub fn from_client_with_activities<ACT>(
         client: impl WorkerClient + 'static,
         act_tasks: ACT,
     ) -> Self
@@ -274,10 +310,9 @@ impl MocksHolder {
         ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
         <ACT as IntoIterator>::IntoIter: Send + 'static,
     {
-        let wft_stream = stream::pending().boxed();
         let mock_act_poller = mock_poller_from_resps(act_tasks);
         let mock_worker = MockWorkerInputs {
-            wft_stream,
+            wft_stream: None,
             act_poller: Some(mock_act_poller),
             nexus_poller: None,
             config: test_worker_cfg().build().unwrap(),
@@ -289,16 +324,78 @@ impl MocksHolder {
         }
     }
 
+    /// Uses the provided list of tasks to create a mock poller with a randomly generated task queue
+    pub fn from_client_with_nexus<NEX>(
+        client: impl WorkerClient + 'static,
+        nexus_tasks: NEX,
+    ) -> Self
+    where
+        NEX: IntoIterator<Item = QueueResponse<PollNexusTaskQueueResponse>>,
+        <NEX as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let mock_nexus_poller = mock_poller_from_resps(nexus_tasks);
+        let mock_worker = MockWorkerInputs {
+            wft_stream: None,
+            act_poller: None,
+            nexus_poller: Some(mock_nexus_poller),
+            config: test_worker_cfg().build().unwrap(),
+        };
+        Self {
+            client: Arc::new(client),
+            inputs: mock_worker,
+            outstanding_task_map: None,
+        }
+    }
+
+    /// Create a MocksHolder with custom combination of task pollers.
+    /// Allows any combination of workflow, activity, and nexus tasks.
+    pub fn from_client_with_custom<WFT, ACT, NEX>(
+        client: impl WorkerClient + 'static,
+        wft_stream: Option<WFT>,
+        activity_tasks: Option<ACT>,
+        nexus_tasks: Option<NEX>,
+    ) -> Self
+    where
+        WFT: Stream<Item = PollWorkflowTaskQueueResponse> + Send + 'static,
+        ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
+        <ACT as IntoIterator>::IntoIter: Send + 'static,
+        NEX: IntoIterator<Item = QueueResponse<PollNexusTaskQueueResponse>>,
+        <NEX as IntoIterator>::IntoIter: Send + 'static,
+    {
+        let wft_stream = wft_stream.map(|s| {
+            s.map(|r| Ok(r.try_into().expect("Mock responses must be valid work")))
+                .boxed()
+        });
+
+        let act_poller = activity_tasks.map(|tasks| mock_poller_from_resps(tasks));
+        let nexus_poller = nexus_tasks.map(|tasks| mock_poller_from_resps(tasks));
+
+        let mock_worker = MockWorkerInputs {
+            wft_stream,
+            act_poller,
+            nexus_poller,
+            config: test_worker_cfg().build().unwrap(),
+        };
+
+        Self {
+            client: Arc::new(client),
+            inputs: mock_worker,
+            outstanding_task_map: None,
+        }
+    }
+
     /// Uses the provided task responses and delivers them as quickly as possible when polled.
     /// This is only useful to test buffering, as typically you do not want to pretend that
     /// the server is delivering WFTs super fast for the same run.
-    pub(crate) fn from_wft_stream(
+    pub fn from_wft_stream(
         client: impl WorkerClient + 'static,
         stream: impl Stream<Item = PollWorkflowTaskQueueResponse> + Send + 'static,
     ) -> Self {
-        let wft_stream = stream
-            .map(|r| Ok(r.try_into().expect("Mock responses must be valid work")))
-            .boxed();
+        let wft_stream = Some(
+            stream
+                .map(|r| Ok(r.try_into().expect("Mock responses must be valid work")))
+                .boxed(),
+        );
         let mock_worker = MockWorkerInputs {
             wft_stream,
             act_poller: None,
@@ -313,8 +410,7 @@ impl MocksHolder {
     }
 }
 
-// TODO: Un-pub ideally
-pub(crate) fn mock_poller_from_resps<T, I>(tasks: I) -> BoxedPoller<T>
+fn mock_poller_from_resps<T, I>(tasks: I) -> BoxedPoller<T>
 where
     T: Send + Sync + 'static,
     I: IntoIterator<Item = QueueResponse<T>>,
@@ -338,16 +434,6 @@ where
     Box::new(mock_poller) as BoxedPoller<T>
 }
 
-pub(crate) fn mock_poller<T>() -> MockPoller<T>
-where
-    T: Send + Sync + 'static,
-{
-    let mut mock_poller = MockPoller::new();
-    mock_poller.expect_shutdown_box().return_const(());
-    mock_poller.expect_notify_shutdown().return_const(());
-    mock_poller
-}
-
 pub(crate) fn mock_manual_poller<T>() -> MockManualPoller<T>
 where
     T: Send + Sync + 'static,
@@ -369,7 +455,7 @@ where
 ///
 /// `num_expected_fails` can be provided to set a specific number of expected failed workflow tasks
 /// sent to the server.
-pub(crate) fn build_multihist_mock_sg(
+pub fn build_multihist_mock_sg(
     hists: impl IntoIterator<Item = FakeWfResponses>,
     enforce_correct_number_of_polls: bool,
     num_expected_fails: usize,
@@ -383,7 +469,7 @@ pub(crate) fn build_multihist_mock_sg(
 }
 
 /// See [build_multihist_mock_sg] -- one history convenience version
-pub(crate) fn single_hist_mock_sg(
+pub fn single_hist_mock_sg(
     wf_id: &str,
     t: TestHistoryBuilder,
     response_batches: impl IntoIterator<Item = impl Into<ResponseType>>,
@@ -399,29 +485,31 @@ type WFTCompletionMockFn = dyn FnMut(&WorkflowTaskCompletion) -> Result<RespondW
     + Send;
 
 #[allow(clippy::type_complexity)]
-pub(crate) struct MockPollCfg {
-    pub(crate) hists: Vec<FakeWfResponses>,
-    pub(crate) enforce_correct_number_of_polls: bool,
-    pub(crate) num_expected_fails: usize,
-    pub(crate) num_expected_legacy_query_resps: usize,
-    pub(crate) mock_client: MockWorkerClient,
+pub struct MockPollCfg {
+    pub hists: Vec<FakeWfResponses>,
+    pub enforce_correct_number_of_polls: bool,
+    pub num_expected_fails: usize,
+    pub num_expected_legacy_query_resps: usize,
+    pub mock_client: MockWorkerClient,
     /// All calls to fail WFTs must match this predicate
-    pub(crate) expect_fail_wft_matcher:
+    pub expect_fail_wft_matcher:
         Box<dyn Fn(&TaskToken, &WorkflowTaskFailedCause, &Option<Failure>) -> bool + Send>,
     /// All calls to legacy query responses must match this predicate
-    pub(crate) expect_legacy_query_matcher:
-        Box<dyn Fn(&TaskToken, &LegacyQueryResult) -> bool + Send>,
-    pub(crate) completion_mock_fn: Option<Box<WFTCompletionMockFn>>,
-    pub(crate) num_expected_completions: Option<TimesRange>,
+    pub expect_legacy_query_matcher: Box<dyn Fn(&TaskToken, &LegacyQueryResult) -> bool + Send>,
+    pub completion_mock_fn: Option<Box<WFTCompletionMockFn>>,
+    pub num_expected_completions: Option<TimesRange>,
     /// If being used with the Rust SDK, this is set true. It ensures pollers will not error out
     /// early with no work, since we cannot know the exact number of times polling will happen.
     /// Instead, they will just block forever.
-    pub(crate) using_rust_sdk: bool,
-    pub(crate) make_poll_stream_interminable: bool,
+    pub using_rust_sdk: bool,
+    pub make_poll_stream_interminable: bool,
+    /// If set, will create a mock activity poller with the provided responeses. If unset, no poller
+    /// is created.
+    pub activity_responses: Option<Vec<QueueResponse<PollActivityTaskQueueResponse>>>,
 }
 
 impl MockPollCfg {
-    pub(crate) fn new(
+    pub fn new(
         hists: Vec<FakeWfResponses>,
         enforce_correct_number_of_polls: bool,
         num_expected_fails: usize,
@@ -438,24 +526,25 @@ impl MockPollCfg {
             num_expected_completions: None,
             using_rust_sdk: false,
             make_poll_stream_interminable: false,
+            activity_responses: None,
         }
     }
 
     /// Builds a config which will hand out each WFT in the history builder one by one
-    pub(crate) fn from_hist_builder(t: TestHistoryBuilder) -> Self {
+    pub fn from_hist_builder(t: TestHistoryBuilder) -> Self {
         let full_hist_info = t.get_full_history_info().unwrap();
         let tasks = 1..=full_hist_info.wf_task_count();
         Self::from_resp_batches("fake_wf_id", t, tasks, mock_worker_client())
     }
 
-    pub(crate) fn from_resps(
+    pub fn from_resps(
         t: TestHistoryBuilder,
         resps: impl IntoIterator<Item = impl Into<ResponseType>>,
     ) -> Self {
         Self::from_resp_batches("fake_wf_id", t, resps, mock_worker_client())
     }
 
-    pub(crate) fn from_resp_batches(
+    pub fn from_resp_batches(
         wf_id: &str,
         t: TestHistoryBuilder,
         resps: impl IntoIterator<Item = impl Into<ResponseType>>,
@@ -477,10 +566,11 @@ impl MockPollCfg {
             num_expected_completions: None,
             using_rust_sdk: false,
             make_poll_stream_interminable: false,
+            activity_responses: None,
         }
     }
 
-    pub(crate) fn completion_asserts_from_expectations(
+    pub fn completion_asserts_from_expectations(
         &mut self,
         builder_fn: impl FnOnce(CompletionAssertsBuilder<'_>),
     ) {
@@ -493,13 +583,13 @@ impl MockPollCfg {
 }
 
 #[allow(clippy::type_complexity)]
-pub(crate) struct CompletionAssertsBuilder<'a> {
+pub struct CompletionAssertsBuilder<'a> {
     dest: &'a mut Option<Box<WFTCompletionMockFn>>,
     assertions: VecDeque<Box<dyn FnOnce(&WorkflowTaskCompletion) + Send>>,
 }
 
 impl CompletionAssertsBuilder<'_> {
-    pub(crate) fn then(
+    pub fn then(
         &mut self,
         assert: impl FnOnce(&WorkflowTaskCompletion) + Send + 'static,
     ) -> &mut Self {
@@ -521,7 +611,7 @@ impl Drop for CompletionAssertsBuilder<'_> {
 }
 
 #[derive(Default, Clone)]
-pub(crate) struct OutstandingWFTMap {
+pub struct OutstandingWFTMap {
     map: Arc<RwLock<BiMap<String, TaskToken>>>,
     waker: Arc<Notify>,
     all_work_delivered: Arc<AtomicBool>,
@@ -541,12 +631,12 @@ impl OutstandingWFTMap {
         self.waker.notify_one();
     }
 
-    pub(crate) fn release_run(&self, run_id: &str) {
+    pub fn release_run(&self, run_id: &str) {
         self.map.write().remove_by_left(run_id);
         self.waker.notify_waiters();
     }
 
-    pub(crate) fn all_work_delivered(&self) -> bool {
+    pub fn all_work_delivered(&self) -> bool {
         self.all_work_delivered.load(Ordering::Acquire)
     }
 }
@@ -571,7 +661,7 @@ impl Drop for EnsuresWorkDoneWFTStream {
 }
 
 /// Given an iterable of fake responses, return the mocks & associated data to work with them
-pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
+pub fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
     let mut task_q_resps: BTreeMap<String, VecDeque<_>> = BTreeMap::new();
     let all_work_delivered = if cfg.enforce_correct_number_of_polls && !cfg.using_rust_sdk {
         Arc::new(AtomicBool::new(false))
@@ -657,7 +747,6 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
                     )
                     .is_err()
                 {
-                    dbg!("Exiting mock WFT task because rcv half of stream was dropped");
                     break;
                 }
             }
@@ -728,15 +817,19 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
         inputs: mock_worker,
         outstanding_task_map: Some(outstanding_wf_task_tokens),
     };
+    if let Some(activity_responses) = cfg.activity_responses {
+        let act_poller = mock_poller_from_resps(activity_responses);
+        mh.set_act_poller(act_poller);
+    }
     if cfg.make_poll_stream_interminable {
         mh.make_wft_stream_interminable();
     }
     mh
 }
 
-pub(crate) struct QueueResponse<T> {
-    pub(crate) resp: T,
-    pub(crate) delay_until: Option<BoxFuture<'static, ()>>,
+pub struct QueueResponse<T> {
+    pub resp: T,
+    pub delay_until: Option<BoxFuture<'static, ()>>,
 }
 
 impl<T> From<T> for QueueResponse<T> {
@@ -777,7 +870,7 @@ where
     }
 }
 
-pub(crate) trait PollWFTRespExt {
+pub trait PollWFTRespExt {
     /// Add an update request to the poll response, using the update name "update_fn" and no args.
     /// Returns the inner request body.
     fn add_update_request(
@@ -820,7 +913,7 @@ impl PollWFTRespExt for PollWorkflowTaskQueueResponse {
     }
 }
 
-pub(crate) fn hist_to_poll_resp(
+pub fn hist_to_poll_resp(
     t: &TestHistoryBuilder,
     wf_id: impl Into<String>,
     response_type: ResponseType,
@@ -857,175 +950,6 @@ pub(crate) fn hist_to_poll_resp(
     QueueResponse { resp, delay_until }
 }
 
-type AsserterWithReply<'a> = (
-    &'a dyn Fn(&WorkflowActivation),
-    workflow_activation_completion::Status,
-);
-
-/// Determines when workflows are kept in the cache or evicted for [poll_and_reply] type tests
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub(crate) enum WorkflowCachingPolicy {
-    /// Workflows are evicted after each workflow task completion. Note that this is *not* after
-    /// each workflow activation - there are often multiple activations per workflow task.
-    NonSticky,
-
-    /// Not a real mode, but good for imitating crashes. Evict workflows after *every* reply,
-    /// even if there are pending activations
-    #[cfg(test)]
-    AfterEveryReply,
-}
-
-/// This function accepts a list of asserts and replies to workflow activations to run against the
-/// provided instance of fake core.
-///
-/// It handles the business of re-sending the same activation replies over again in the event
-/// of eviction or workflow activation failure. Activation failures specifically are only run once,
-/// since they clearly can't be returned every time we replay the workflow, or it could never
-/// proceed
-pub(crate) async fn poll_and_reply<'a>(
-    worker: &'a Worker,
-    eviction_mode: WorkflowCachingPolicy,
-    expect_and_reply: &'a [AsserterWithReply<'a>],
-) {
-    poll_and_reply_clears_outstanding_evicts(worker, None, eviction_mode, expect_and_reply).await;
-}
-
-pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
-    worker: &'a Worker,
-    outstanding_map: Option<OutstandingWFTMap>,
-    eviction_mode: WorkflowCachingPolicy,
-    expect_and_reply: &'a [AsserterWithReply<'a>],
-) {
-    let mut evictions = 0;
-    let expected_evictions = expect_and_reply.len() - 1;
-    let mut executed_failures = HashSet::new();
-    let expected_fail_count = expect_and_reply
-        .iter()
-        .filter(|(_, reply)| !reply.is_success())
-        .count();
-
-    'outer: loop {
-        let expect_iter = expect_and_reply.iter();
-
-        for (i, interaction) in expect_iter.enumerate() {
-            let (asserter, reply) = interaction;
-            let complete_is_failure = !reply.is_success();
-            // Only send activation failures once
-            if executed_failures.contains(&i) {
-                continue;
-            }
-
-            let mut res = worker.poll_workflow_activation().await.unwrap();
-            if res.jobs.iter().any(|j| {
-                matches!(
-                    j.variant,
-                    Some(workflow_activation_job::Variant::RemoveFromCache(_))
-                )
-            }) && res.jobs.len() > 1
-            {
-                panic!("Saw an activation with an eviction & other work! {res:?}");
-            }
-            let is_eviction = res.is_only_eviction();
-
-            let mut do_release = false;
-
-            if is_eviction {
-                // If the job is an eviction, clear it, since in the tests we don't explicitly
-                // specify evict assertions
-                res.jobs.clear();
-                do_release = true;
-            }
-
-            // TODO: Can remove this if?
-            if !res.jobs.is_empty() {
-                asserter(&res);
-            }
-
-            let reply = if res.jobs.is_empty() {
-                // Just an eviction
-                WorkflowActivationCompletion::empty(res.run_id.clone())
-            } else {
-                // Eviction plus some work, we still want to issue the reply
-                WorkflowActivationCompletion {
-                    run_id: res.run_id.clone(),
-                    status: Some(reply.clone()),
-                }
-            };
-
-            let ends_execution = reply.has_execution_ending();
-
-            worker.complete_workflow_activation(reply).await.unwrap();
-
-            if do_release && let Some(omap) = outstanding_map.as_ref() {
-                omap.release_run(&res.run_id);
-            }
-            // Restart assertions from the beginning if it was an eviction (and workflow execution
-            // isn't over)
-            if is_eviction && !ends_execution {
-                continue 'outer;
-            }
-
-            if complete_is_failure {
-                executed_failures.insert(i);
-            }
-
-            match eviction_mode {
-                WorkflowCachingPolicy::NonSticky => (),
-                WorkflowCachingPolicy::AfterEveryReply => {
-                    if evictions < expected_evictions {
-                        worker.request_workflow_eviction(&res.run_id);
-                        evictions += 1;
-                    }
-                }
-            }
-        }
-
-        break;
-    }
-
-    assert_eq!(expected_fail_count, executed_failures.len());
-    assert_eq!(worker.outstanding_workflow_tasks().await, 0);
-}
-
-pub(crate) fn gen_assert_and_reply(
-    asserter: &dyn Fn(&WorkflowActivation),
-    reply_commands: Vec<workflow_command::Variant>,
-) -> AsserterWithReply<'_> {
-    (
-        asserter,
-        workflow_completion::Success::from_variants(reply_commands).into(),
-    )
-}
-
-pub(crate) fn gen_assert_and_fail(asserter: &dyn Fn(&WorkflowActivation)) -> AsserterWithReply<'_> {
-    (
-        asserter,
-        workflow_completion::Failure {
-            failure: Some(Failure {
-                message: "Intentional test failure".to_string(),
-                ..Default::default()
-            }),
-            ..Default::default()
-        }
-        .into(),
-    )
-}
-
-/// Generate asserts for [poll_and_reply] by passing patterns to match against the job list
-#[macro_export]
-macro_rules! job_assert {
-    ($($pat:pat),+) => {
-        |res| {
-            assert_matches!(
-                res.jobs.as_slice(),
-                [$(WorkflowActivationJob {
-                    variant: Some($pat),
-                }),+]
-            );
-        }
-    };
-}
-
 /// Forcibly drive a future a number of times, enforcing it is always returning Pending. This is
 /// useful for ensuring some future has proceeded "enough" before racing it against another future.
 #[macro_export]
@@ -1054,7 +978,7 @@ macro_rules! prost_dur {
 }
 
 #[async_trait]
-pub(crate) trait WorkerExt {
+pub trait WorkerExt {
     /// Initiate shutdown, drain the pollers, and wait for shutdown to complete.
     async fn drain_pollers_and_shutdown(self);
     /// Initiate shutdown, drain the *activity* poller, and wait for shutdown to complete.
@@ -1065,30 +989,7 @@ pub(crate) trait WorkerExt {
 #[async_trait]
 impl WorkerExt for Worker {
     async fn drain_pollers_and_shutdown(self) {
-        self.initiate_shutdown();
-        tokio::join!(
-            async {
-                assert_matches!(
-                    self.poll_activity_task().await.unwrap_err(),
-                    PollError::ShutDown
-                );
-            },
-            async {
-                loop {
-                    match self.poll_workflow_activation().await {
-                        Err(PollError::ShutDown) => break,
-                        Ok(a) if a.is_only_eviction() => {
-                            self.complete_workflow_activation(WorkflowActivationCompletion::empty(
-                                a.run_id,
-                            ))
-                            .await
-                            .unwrap();
-                        }
-                        o => panic!("Unexpected activation while draining: {o:?}"),
-                    }
-                }
-            }
-        );
+        drain_pollers_and_shutdown(&self).await;
         self.finalize_shutdown().await;
     }
 
@@ -1101,3 +1002,57 @@ impl WorkerExt for Worker {
         self.shutdown().await;
     }
 }
+
+#[async_trait::async_trait]
+/// Test helper methods for core workers
+pub trait WorkerTestHelpers {
+    /// Complete a workflow execution
+    async fn complete_execution(&self, run_id: &str);
+    /// Complete a timer with the given sequence number and duration
+    async fn complete_timer(&self, run_id: &str, seq: u32, duration: Duration);
+    /// Handle workflow eviction from cache
+    async fn handle_eviction(&self);
+}
+
+#[async_trait::async_trait]
+impl<T> WorkerTestHelpers for T
+where
+    T: WorkerTrait + ?Sized,
+{
+    async fn complete_execution(&self, run_id: &str) {
+        self.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+            run_id.to_string(),
+            vec![CompleteWorkflowExecution { result: None }.into()],
+        ))
+        .await
+        .unwrap();
+    }
+
+    async fn complete_timer(&self, run_id: &str, seq: u32, duration: Duration) {
+        self.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+            run_id.to_string(),
+            vec![
+                StartTimer {
+                    seq,
+                    start_to_fire_timeout: Some(duration.try_into().expect("duration fits")),
+                }
+                .into(),
+            ],
+        ))
+        .await
+        .unwrap();
+    }
+
+    async fn handle_eviction(&self) {
+        let task = self.poll_workflow_activation().await.unwrap();
+        assert_matches!(
+            task.jobs.as_slice(),
+            [WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::RemoveFromCache(_)),
+            }]
+        );
+        self.complete_workflow_activation(WorkflowActivationCompletion::empty(task.run_id))
+            .await
+            .unwrap();
+    }
+}