@temporalio/core-bridge 0.19.2 → 0.20.2

package/sdk-core/core/src/worker/mod.rs

@@ -1,5 +1,5 @@
 mod activities;
-mod dispatcher;
+pub(crate) mod client;
 mod wft_delivery;
 
 pub use temporal_sdk_core_api::worker::{WorkerConfig, WorkerConfigBuilder};
@@ -8,20 +8,25 @@ pub(crate) use activities::{
     ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
     NewLocalAct,
 };
-pub(crate) use dispatcher::WorkerDispatcher;
 
 use crate::{
+    abstractions::MeteredSemaphore,
     errors::CompleteWfError,
     pollers::{
         new_activity_task_buffer, new_workflow_task_buffer, BoxedActPoller, BoxedWFPoller, Poller,
         WorkflowTaskPoller,
     },
     protosext::{legacy_query_failure, ValidPollWFTQResponse},
-    telemetry::metrics::{
-        activity_poller, workflow_poller, workflow_sticky_poller, MetricsContext,
+    telemetry::{
+        metrics::{
+            activity_poller, local_activity_worker_type, workflow_poller, workflow_sticky_poller,
+            workflow_worker_type, MetricsContext,
+        },
+        VecDisplayer,
     },
     worker::{
         activities::{DispatchOrTimeoutLA, LACompleteAction, LocalActivityManager},
+        client::WorkerClientBag,
         wft_delivery::WFTSource,
     },
     workflow::{
@@ -31,18 +36,19 @@ use crate::{
         },
         EmptyWorkflowCommandErr, LocalResolution, WFMachinesError, WorkflowCachingPolicy,
     },
-    ActivityHeartbeat, CompleteActivityError, PollActivityError, PollWfError,
+    ActivityHeartbeat, CompleteActivityError, PollActivityError, PollWfError, WorkerTrait,
 };
 use activities::{LocalInFlightActInfo, WorkerActivityTasks};
 use futures::{Future, TryFutureExt};
 use std::{convert::TryInto, sync::Arc};
-use temporal_client::{ServerGatewayApis, WorkflowTaskCompletion};
+use temporal_client::WorkflowTaskCompletion;
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::activity_execution_result,
         activity_task::ActivityTask,
         workflow_activation::{remove_from_cache::EvictionReason, WorkflowActivation},
         workflow_completion::{self, workflow_activation_completion, WorkflowActivationCompletion},
+        ActivityTaskCompletion,
     },
     temporal::api::{
         enums::v1::{TaskQueueKind, WorkflowTaskFailedCause},
@@ -52,14 +58,19 @@ use temporal_sdk_core_protos::{
     },
     TaskToken,
 };
-use tokio::sync::{watch, Notify, Semaphore};
+use tokio::sync::Notify;
+use tokio_util::sync::CancellationToken;
 use tonic::Code;
 use tracing_futures::Instrument;
 
+#[cfg(test)]
+use crate::worker::client::WorkerClient;
+use crate::workflow::workflow_tasks::EvictionRequestResult;
+
 /// A worker polls on a certain task queue
 pub struct Worker {
     config: WorkerConfig,
-    server_gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
+    wf_client: Arc<WorkerClientBag>,
 
     /// Will be populated when this worker should poll on a sticky WFT queue
     sticky_name: Option<String>,
@@ -67,14 +78,14 @@ pub struct Worker {
     /// Buffers workflow task polling in the event we need to return a pending activation while
     /// a poll is ongoing. Sticky and nonsticky polling happens inside of it.
     wf_task_source: WFTSource,
-    /// Workflow task management TODO: No pub
-    pub(crate) wft_manager: WorkflowTaskManager,
+    /// Workflow task management
+    wft_manager: WorkflowTaskManager,
     /// Manages activity tasks for this worker/task queue
     at_task_mgr: Option<WorkerActivityTasks>,
     /// Manages local activities
     local_act_mgr: LocalActivityManager,
     /// Ensures we stay at or below this worker's maximum concurrent workflow limit
-    workflows_semaphore: Semaphore,
+    workflows_semaphore: MeteredSemaphore,
     /// Used to wake blocked workflow task polling when there is some change to workflow activations
     /// that should cause us to restart the loop
     pending_activations_notify: Arc<Notify>,
@@ -82,21 +93,106 @@ pub struct Worker {
     /// a WFT is completed.
     wfts_drained_notify: Arc<Notify>,
     /// Has shutdown been called?
-    shutdown_requested: watch::Receiver<bool>,
-    shutdown_sender: watch::Sender<bool>,
+    shutdown_token: CancellationToken,
     /// Will be called at the end of each activation completion
     post_activate_hook: Option<Box<dyn Fn(&Self) + Send + Sync>>,
 
     metrics: MetricsContext,
 }
 
+#[async_trait::async_trait]
+impl WorkerTrait for Worker {
+    #[instrument(level = "debug", skip(self), fields(run_id))]
+    async fn poll_workflow_activation(&self) -> Result<WorkflowActivation, PollWfError> {
+        self.next_workflow_activation().await
+    }
+
+    #[instrument(level = "debug", skip(self))]
+    async fn poll_activity_task(&self) -> Result<ActivityTask, PollActivityError> {
+        loop {
+            match self.activity_poll().await.transpose() {
+                Some(r) => break r,
+                None => {
+                    tokio::task::yield_now().await;
+                    continue;
+                }
+            }
+        }
+    }
+
+    #[instrument(level = "debug", skip(self, completion),
+    fields(completion=%&completion, run_id=%completion.run_id))]
+    async fn complete_workflow_activation(
+        &self,
+        completion: WorkflowActivationCompletion,
+    ) -> Result<(), CompleteWfError> {
+        self.complete_workflow_activation(completion).await
+    }
+
+    #[instrument(level = "debug", skip(self, completion),
+    fields(completion=%&completion))]
+    async fn complete_activity_task(
+        &self,
+        completion: ActivityTaskCompletion,
+    ) -> Result<(), CompleteActivityError> {
+        let task_token = TaskToken(completion.task_token);
+        let status = if let Some(s) = completion.result.and_then(|r| r.status) {
+            s
+        } else {
+            return Err(CompleteActivityError::MalformedActivityCompletion {
+                reason: "Activity completion had empty result/status field".to_owned(),
+                completion: None,
+            });
+        };
+
+        self.complete_activity(task_token, status).await
+    }
+
+    fn record_activity_heartbeat(&self, details: ActivityHeartbeat) {
+        self.record_heartbeat(details);
+    }
+
+    fn request_workflow_eviction(&self, run_id: &str) {
+        self.request_wf_eviction(
+            run_id,
+            "Eviction explicitly requested by lang",
+            EvictionReason::LangRequested,
+        );
+    }
+
+    fn get_config(&self) -> &WorkerConfig {
+        &self.config
+    }
+
+    /// Begins the shutdown process, tells pollers they should stop. Is idempotent.
+    // TODO: will be in trait after Roey's shutdown refactor
+    fn initiate_shutdown(&self) {
+        self.shutdown_token.cancel();
+        // First, we want to stop polling of both activity and workflow tasks
+        if let Some(atm) = self.at_task_mgr.as_ref() {
+            atm.notify_shutdown();
+        }
+        self.wf_task_source.stop_pollers();
+    }
+
+    async fn shutdown(&self) {
+        self.shutdown().await
+    }
+
+    async fn finalize_shutdown(self) {
+        self.shutdown().await;
+        self.finalize_shutdown().await
+    }
+}
+
 impl Worker {
     pub(crate) fn new(
         config: WorkerConfig,
         sticky_queue_name: Option<String>,
-        sg: Arc<dyn ServerGatewayApis + Send + Sync>,
+        client: Arc<WorkerClientBag>,
         metrics: MetricsContext,
     ) -> Self {
+        info!(task_queue = %config.task_queue, "Initializing worker");
         metrics.worker_registered();
 
         let max_nonsticky_polls = if sticky_queue_name.is_some() {
@@ -107,7 +203,7 @@ impl Worker {
         let max_sticky_polls = config.max_sticky_polls();
         let wft_metrics = metrics.with_new_attrs([workflow_poller()]);
         let mut wf_task_poll_buffer = new_workflow_task_buffer(
-            sg.clone(),
+            client.clone(),
             config.task_queue.clone(),
             false,
             max_nonsticky_polls,
@@ -117,7 +213,7 @@ impl Worker {
         let sticky_queue_poller = sticky_queue_name.as_ref().map(|sqn| {
             let sticky_metrics = metrics.with_new_attrs([workflow_sticky_poller()]);
             let mut sp = new_workflow_task_buffer(
-                sg.clone(),
+                client.clone(),
                 sqn.clone(),
                 true,
                 max_sticky_polls,
@@ -130,10 +226,11 @@ impl Worker {
             None
         } else {
             let mut ap = new_activity_task_buffer(
-                sg.clone(),
+                client.clone(),
                 config.task_queue.clone(),
                 config.max_concurrent_at_polls,
                 config.max_concurrent_at_polls * 2,
+                config.max_task_queue_activities_per_second,
             );
             let act_metrics = metrics.with_new_attrs([activity_poller()]);
             ap.set_num_pollers_handler(move |np| act_metrics.record_num_pollers(np));
@@ -149,17 +246,22 @@ impl Worker {
         Self::new_with_pollers(
             config,
             sticky_queue_name,
-            sg,
+            client,
             wf_task_poll_buffer,
             act_poll_buffer,
             metrics,
         )
     }
 
+    #[cfg(test)]
+    pub(crate) fn new_test(config: WorkerConfig, client: impl WorkerClient + 'static) -> Self {
+        Self::new(config, None, Arc::new(client.into()), Default::default())
+    }
+
     pub(crate) fn new_with_pollers(
         config: WorkerConfig,
         sticky_queue_name: Option<String>,
-        sg: Arc<dyn ServerGatewayApis + Send + Sync>,
+        client: Arc<WorkerClientBag>,
         wft_poller: BoxedWFPoller,
         act_poller: Option<BoxedActPoller>,
         metrics: MetricsContext,
@@ -173,9 +275,8 @@ impl Worker {
         };
         let pa_notif = Arc::new(Notify::new());
         let wfts_drained_notify = Arc::new(Notify::new());
-        let (shut_tx, shut_rx) = watch::channel(false);
         Self {
-            server_gateway: sg.clone(),
+            wf_client: client.clone(),
             sticky_name: sticky_queue_name,
             wf_task_source: WFTSource::new(wft_poller),
             wft_manager: WorkflowTaskManager::new(pa_notif.clone(), cache_policy, metrics.clone()),
@@ -183,7 +284,7 @@ impl Worker {
                 WorkerActivityTasks::new(
                     config.max_outstanding_activities,
                     ap,
-                    sg.clone(),
+                    client.clone(),
                     metrics.clone(),
                     config.max_heartbeat_throttle_interval,
                     config.default_heartbeat_throttle_interval,
@@ -191,12 +292,16 @@ impl Worker {
             }),
             local_act_mgr: LocalActivityManager::new(
                 config.max_outstanding_local_activities,
-                sg.get_options().namespace.clone(),
+                config.namespace.clone(),
+                metrics.with_new_attrs([local_activity_worker_type()]),
+            ),
+            workflows_semaphore: MeteredSemaphore::new(
+                config.max_outstanding_workflow_tasks,
+                metrics.with_new_attrs([workflow_worker_type()]),
+                MetricsContext::available_task_slots,
             ),
-            workflows_semaphore: Semaphore::new(config.max_outstanding_workflow_tasks),
             config,
-            shutdown_requested: shut_rx,
-            shutdown_sender: shut_tx,
+            shutdown_token: CancellationToken::new(),
             post_activate_hook: None,
             pending_activations_notify: pa_notif,
             wfts_drained_notify,
@@ -204,20 +309,11 @@ impl Worker {
         }
     }
 
-    /// Begins the shutdown process, tells pollers they should stop. Is idempotent.
-    pub(crate) fn initiate_shutdown(&self) {
-        let _ = self.shutdown_sender.send(true);
-        // First, we want to stop polling of both activity and workflow tasks
-        if let Some(atm) = self.at_task_mgr.as_ref() {
-            atm.notify_shutdown();
-        }
-        self.wf_task_source.stop_pollers();
-    }
-
     /// Will shutdown the worker. Does not resolve until all outstanding workflow tasks have been
     /// completed
     pub(crate) async fn shutdown(&self) {
         self.initiate_shutdown();
+        info!("Initiated shutdown");
         // Next we need to wait for all local activities to finish so no more workflow task
         // heartbeats will be generated
         self.local_act_mgr.shutdown_and_wait_all_finished().await;
@@ -249,7 +345,7 @@ impl Worker {
 
     #[cfg(test)]
     pub(crate) fn available_wft_permits(&self) -> usize {
-        self.workflows_semaphore.available_permits()
+        self.workflows_semaphore.sem.available_permits()
     }
 
     /// Get new activity tasks (may be local or nonlocal). Local activities are returned first
@@ -257,12 +353,12 @@ impl Worker {
     ///
     /// Returns `Ok(None)` in the event of a poll timeout or if the polling loop should otherwise
     /// be restarted
-    pub(crate) async fn activity_poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
+    async fn activity_poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
         let act_mgr_poll = async {
             if let Some(ref act_mgr) = self.at_task_mgr {
                 act_mgr.poll().await
             } else {
-                let _ = self.shutdown_requested.clone().changed().await;
+                self.shutdown_token.cancelled().await;
                 Err(PollActivityError::ShutDown)
             }
         };
@@ -278,7 +374,12 @@ impl Worker {
                             &run_id, LocalResolution::LocalActivity(resolution)).await;
                         Ok(task)
                     },
-                    None => Ok(None)
+                    None => {
+                        if self.shutdown_token.is_cancelled() {
+                            return Err(PollActivityError::ShutDown);
+                        }
+                        Ok(None)
+                    }
                 }
             },
             r = act_mgr_poll => r,
@@ -326,8 +427,7 @@ impl Worker {
         }
 
         if let Some(atm) = &self.at_task_mgr {
-            atm.complete(task_token, status, self.server_gateway.as_ref())
-                .await
+            atm.complete(task_token, status, &**self.wf_client).await
         } else {
             error!(
                 "Tried to complete activity {} on a worker that does not have an activity manager",
@@ -428,16 +528,21 @@ impl Worker {
 
     /// Tell the worker a workflow task has completed, for tracking max outstanding WFTs
     pub(crate) fn return_workflow_task_permit(&self) {
-        self.workflows_semaphore.add_permits(1);
+        self.workflows_semaphore.add_permit();
     }
 
+    /// Request a workflow eviction. Returns true if we actually queued up a new eviction request.
     pub(crate) fn request_wf_eviction(
         &self,
         run_id: &str,
         message: impl Into<String>,
         reason: EvictionReason,
-    ) {
-        self.wft_manager.request_eviction(run_id, message, reason);
+    ) -> bool {
+        match self.wft_manager.request_eviction(run_id, message, reason) {
+            EvictionRequestResult::EvictionIssued(_) => true,
+            EvictionRequestResult::NotFound => false,
+            EvictionRequestResult::EvictionAlreadyOutstanding => false,
+        }
     }
 
     /// Sets a function to be called at the end of each activation completion
@@ -448,11 +553,25 @@ impl Worker {
         self.post_activate_hook = Some(Box::new(callback))
     }
 
+    /// Used for replay workers - causes the worker to shutdown when the given run reaches the
+    /// given event number
+    pub(crate) fn set_shutdown_on_run_reaches_event(&mut self, run_id: String, last_event: i64) {
+        self.set_post_activate_hook(move |worker| {
+            if worker
+                .wft_manager
+                .most_recently_processed_event(&run_id)
+                .unwrap_or_default()
+                >= last_event
+            {
+                worker.initiate_shutdown();
+            }
+        });
+    }
+
     /// Resolves with WFT poll response or `PollWfError::ShutDown` if WFTs have been drained
     async fn workflow_poll_or_wfts_drained(
         &self,
     ) -> Result<Option<ValidPollWFTQResponse>, PollWfError> {
-        let mut shutdown_requested = self.shutdown_requested.clone();
         loop {
             tokio::select! {
                 biased;
@@ -466,7 +585,7 @@ impl Worker {
                     }
                     return r
                 },
-                _ = shutdown_requested.changed() => {},
+                _ = self.shutdown_token.cancelled() => {},
             }
         }
     }
@@ -481,7 +600,7 @@ impl Worker {
         // heartbeating which is a "new" workflow task that we need to accept and process as long as
         // the LA is outstanding. Similarly, if we already have such tasks (from a WFT completion),
         // then we must fetch them from the source before we can say workflow polling is shutdown.
-        if *self.shutdown_requested.borrow()
+        if self.shutdown_token.is_cancelled()
             && !self.wf_task_source.has_tasks_from_complete()
             && self.local_act_mgr.num_outstanding() == 0
         {
@@ -535,10 +654,9 @@ impl Worker {
         work: ValidPollWFTQResponse,
     ) -> Result<Option<WorkflowActivation>, PollWfError> {
         let we = work.workflow_execution.clone();
-        let tt = work.task_token.clone();
         let res = self
             .wft_manager
-            .apply_new_poll_resp(work, self.server_gateway.clone())
+            .apply_new_poll_resp(work, self.wf_client.clone())
             .await;
         Ok(match res {
             NewWfTaskOutcome::IssueActivation(a) => {
@@ -555,38 +673,24 @@ impl Worker {
                 debug!(workflow_execution=?we,
                        "No new work for lang to perform after polling server");
                 self.complete_workflow_activation(WorkflowActivationCompletion {
-                    task_queue: self.config.task_queue.clone(),
                     run_id: we.run_id,
                     status: Some(workflow_completion::Success::from_variants(vec![]).into()),
                 })
                 .await?;
                 None
             }
-            NewWfTaskOutcome::CacheMiss => {
-                debug!(workflow_execution=?we, "Unable to process workflow task with partial \
-                history because workflow cache does not contain workflow anymore.");
-                self.server_gateway
-                    .fail_workflow_task(
-                        tt,
-                        WorkflowTaskFailedCause::ResetStickyTaskQueue,
-                        Some(Failure {
-                            message: "Unable to process workflow task with partial history \
-                                      because workflow cache does not contain workflow anymore."
-                                .to_string(),
-                            ..Default::default()
-                        }),
-                    )
-                    .await?;
-                self.return_workflow_task_permit();
-                None
-            }
             NewWfTaskOutcome::Evict(e) => {
                 warn!(error=?e, run_id=%we.run_id, "Error while applying poll response to workflow");
-                self.request_wf_eviction(
+                let did_issue_eviction = self.request_wf_eviction(
                     &we.run_id,
                     format!("Error while applying poll response to workflow: {:?}", e),
                     e.evict_reason(),
                 );
+                // If we didn't actually need to issue an eviction, then return the WFT permit.
+                // EX: The workflow we tried to evict wasn't in the cache.
+                if !did_issue_eviction {
+                    self.return_workflow_task_permit();
+                }
                 None
             }
         })
@@ -627,9 +731,12 @@ impl Worker {
                         force_new_wft,
                     },
             })) => {
-                debug!("Sending commands to server: {:?}", &commands);
+                debug!("Sending commands to server: {}", commands.display());
                 if !query_responses.is_empty() {
-                    debug!("Sending query responses to server: {:?}", &query_responses);
+                    debug!(
+                        "Sending query responses to server: {}",
+                        query_responses.display()
+                    );
                 }
                 let mut completion = WorkflowTaskCompletion {
                     task_token,
@@ -649,7 +756,7 @@ impl Worker {
 
                 self.handle_wft_reporting_errs(run_id, || async {
                     let maybe_wft = self
-                        .server_gateway
+                        .wf_client
                         .complete_workflow_task(completion)
                         .instrument(span!(tracing::Level::DEBUG, "Complete WFT call"))
                         .await?;
@@ -669,7 +776,7 @@ impl Worker {
                 action: ActivationAction::RespondLegacyQuery { result },
                 ..
             })) => {
-                self.server_gateway
+                self.wf_client
                     .respond_legacy_query(task_token, result)
                     .await?;
                 Ok(WFTReportOutcome {
@@ -720,7 +827,7 @@ impl Worker {
                 FailedActivationOutcome::Report(tt) => {
                     warn!(run_id, failure=?failure, "Failing workflow activation");
                     self.handle_wft_reporting_errs(run_id, || async {
-                        self.server_gateway
+                        self.wf_client
                             .fail_workflow_task(tt, cause, failure.failure.map(Into::into))
                             .await
                     })
@@ -732,7 +839,7 @@ impl Worker {
                 }
                 FailedActivationOutcome::ReportLegacyQueryFailure(task_token) => {
                     warn!(run_id, failure=?failure, "Failing legacy query request");
-                    self.server_gateway
+                    self.wf_client
                         .respond_legacy_query(task_token, legacy_query_failure(failure))
                         .await?;
                     WFTReportOutcome {
@@ -765,7 +872,7 @@ impl Worker {
                     // Silence unhandled command errors since the lang SDK cannot do anything about
                     // them besides poll again, which it will do anyway.
                     tonic::Code::InvalidArgument if err.message() == "UnhandledCommand" => {
-                        warn!(error = %err, run_id, "Unhandled command response when completing");
+                        debug!(error = %err, run_id, "Unhandled command response when completing");
                         should_evict = Some(EvictionReason::UnhandledCommand);
                         Ok(())
                     }
@@ -851,91 +958,83 @@ struct WFTReportOutcome {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use temporal_client::mocks::mock_gateway;
+    use crate::{test_help::test_worker_cfg, worker::client::mocks::mock_workflow_client};
     use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse;
 
     #[tokio::test]
     async fn activity_timeouts_dont_eat_permits() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_poll_activity_task()
-            .returning(|_| Ok(PollActivityTaskQueueResponse::default()));
+            .returning(|_, _| Ok(PollActivityTaskQueueResponse::default()));
 
-        let cfg = WorkerConfigBuilder::default()
-            .task_queue("whatever")
+        let cfg = test_worker_cfg()
             .max_outstanding_activities(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
+        let worker = Worker::new_test(cfg, mock_client);
         assert_eq!(worker.activity_poll().await.unwrap(), None);
         assert_eq!(worker.at_task_mgr.unwrap().remaining_activity_capacity(), 5);
     }
 
     #[tokio::test]
     async fn workflow_timeouts_dont_eat_permits() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_poll_workflow_task()
             .returning(|_, _| Ok(PollWorkflowTaskQueueResponse::default()));
 
-        let cfg = WorkerConfigBuilder::default()
-            .task_queue("whatever")
+        let cfg = test_worker_cfg()
             .max_outstanding_workflow_tasks(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
+        let worker = Worker::new_test(cfg, mock_client);
         assert_eq!(worker.workflow_poll().await.unwrap(), None);
-        assert_eq!(worker.workflows_semaphore.available_permits(), 5);
+        assert_eq!(worker.workflows_semaphore.sem.available_permits(), 5);
     }
 
     #[tokio::test]
     async fn activity_errs_dont_eat_permits() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_poll_activity_task()
-            .returning(|_| Err(tonic::Status::internal("ahhh")));
+            .returning(|_, _| Err(tonic::Status::internal("ahhh")));
 
-        let cfg = WorkerConfigBuilder::default()
-            .task_queue("whatever")
+        let cfg = test_worker_cfg()
             .max_outstanding_activities(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
+        let worker = Worker::new_test(cfg, mock_client);
         assert!(worker.activity_poll().await.is_err());
         assert_eq!(worker.at_task_mgr.unwrap().remaining_activity_capacity(), 5);
     }
 
     #[tokio::test]
     async fn workflow_errs_dont_eat_permits() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_poll_workflow_task()
             .returning(|_, _| Err(tonic::Status::internal("ahhh")));
 
-        let cfg = WorkerConfigBuilder::default()
-            .task_queue("whatever")
+        let cfg = test_worker_cfg()
             .max_outstanding_workflow_tasks(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
+        let worker = Worker::new_test(cfg, mock_client);
         assert!(worker.workflow_poll().await.is_err());
-        assert_eq!(worker.workflows_semaphore.available_permits(), 5);
+        assert_eq!(worker.workflows_semaphore.sem.available_permits(), 5);
     }
 
     #[test]
     fn max_polls_calculated_properly() {
-        let cfg = WorkerConfigBuilder::default()
-            .task_queue("whatever")
-            .build()
-            .unwrap();
+        let cfg = test_worker_cfg().build().unwrap();
         assert_eq!(cfg.max_nonsticky_polls(), 1);
         assert_eq!(cfg.max_sticky_polls(), 4);
     }
 
     #[test]
     fn max_polls_zero_is_err() {
-        assert!(WorkerConfigBuilder::default()
-            .task_queue("whatever")
+        assert!(test_worker_cfg()
             .max_concurrent_wft_polls(0_usize)
             .build()
             .is_err());