@temporalio/core-bridge 0.23.0 → 1.0.0

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

@@ -2,22 +2,31 @@ pub(crate) use temporal_sdk_core_test_utils::canned_histories;
 
 use crate::{
     pollers::{BoxedActPoller, BoxedPoller, BoxedWFPoller, MockManualPoller, MockPoller},
+    protosext::ValidPollWFTQResponse,
     replay::TestHistoryBuilder,
     sticky_q_name_for_worker,
-    worker::client::{mocks::mock_workflow_client, MockWorkerClient, WorkerClient},
-    workflow::WorkflowCachingPolicy,
+    worker::{
+        client::{mocks::mock_workflow_client, MockWorkerClient, WorkerClient},
+        new_wft_poller,
+    },
     TaskToken, Worker, WorkerClientBag, WorkerConfig, WorkerConfigBuilder,
 };
 use bimap::BiMap;
-use futures::FutureExt;
+use futures::{future::BoxFuture, stream, stream::BoxStream, FutureExt, Stream, StreamExt};
 use mockall::TimesRange;
 use parking_lot::RwLock;
 use std::{
     collections::{BTreeMap, HashMap, HashSet, VecDeque},
-    ops::RangeFull,
-    sync::Arc,
+    ops::{Deref, DerefMut},
+    pin::Pin,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    task::{Context, Poll},
     time::Duration,
 };
+use temporal_client::WorkflowTaskCompletion;
 use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::{
     coresdk::{
@@ -36,39 +45,61 @@ use temporal_sdk_core_protos::{
     },
 };
 use temporal_sdk_core_test_utils::TestWorker;
-use tokio::sync::Notify;
+use tokio::sync::{mpsc::unbounded_channel, Notify};
+use tokio_stream::wrappers::UnboundedReceiverStream;
+use tokio_util::sync::CancellationToken;
 
 pub const TEST_Q: &str = "q";
 pub static NO_MORE_WORK_ERROR_MSG: &str = "No more work to do";
 
 pub fn test_worker_cfg() -> WorkerConfigBuilder {
     let mut wcb = WorkerConfigBuilder::default();
-    wcb.namespace("default").task_queue(TEST_Q);
+    wcb.namespace("default")
+        .task_queue(TEST_Q)
+        .worker_build_id("test_bin_id")
+        // Serial polling since it makes mocking much easier.
+        .max_concurrent_wft_polls(1_usize);
     wcb
 }
 
 /// When constructing responses for mocks, indicates how a given response should be built
-#[derive(derive_more::From, Debug, Clone, Copy, Eq, PartialEq, Hash)]
+#[derive(derive_more::From)]
+#[allow(clippy::large_enum_variant)] // Test only code, whatever.
 pub enum ResponseType {
     ToTaskNum(usize),
     /// Returns just the history after 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
     #[from(ignore)]
     OneTask(usize),
+    /// Waits until the future resolves before responding as `ToTaskNum` with the provided number
+    UntilResolved(BoxFuture<'static, ()>, usize),
     AllHistory,
+    Raw(PollWorkflowTaskQueueResponse),
+}
+#[derive(Eq, PartialEq, Hash)]
+pub enum HashableResponseType {
+    ToTaskNum(usize),
+    OneTask(usize),
+    UntilResolved(usize),
+    AllHistory,
+    Raw(TaskToken),
+}
+impl ResponseType {
+    pub fn hashable(&self) -> HashableResponseType {
+        match self {
+            ResponseType::ToTaskNum(x) => HashableResponseType::ToTaskNum(*x),
+            ResponseType::OneTask(x) => HashableResponseType::OneTask(*x),
+            ResponseType::AllHistory => HashableResponseType::AllHistory,
+            ResponseType::Raw(r) => HashableResponseType::Raw(r.task_token.clone().into()),
+            ResponseType::UntilResolved(_, x) => HashableResponseType::UntilResolved(*x),
+        }
+    }
 }
-
 impl From<&usize> for ResponseType {
     fn from(u: &usize) -> Self {
         Self::ToTaskNum(*u)
     }
 }
-// :shrug:
-impl From<&Self> for ResponseType {
-    fn from(r: &Self) -> Self {
-        *r
-    }
-}
 
 /// Given identifiers for a workflow/run, and a test history builder, construct an instance of
 /// the a worker with a mock server client that will produce the responses as appropriate.
@@ -89,20 +120,21 @@ pub(crate) fn build_fake_worker(
             response_batches,
         }],
         true,
-        None,
+        0,
     );
     mock_worker(mock_holder)
 }
 
 pub(crate) fn mock_worker(mocks: MocksHolder) -> Worker {
-    let sticky_q = sticky_q_name_for_worker("unit-test", &mocks.mock_worker.config);
+    let sticky_q = sticky_q_name_for_worker("unit-test", &mocks.inputs.config);
     Worker::new_with_pollers(
-        mocks.mock_worker.config,
+        mocks.inputs.config,
         sticky_q,
         Arc::new(mocks.client_bag),
-        mocks.mock_worker.wf_poller,
-        mocks.mock_worker.act_poller,
+        mocks.inputs.wft_stream,
+        mocks.inputs.act_poller,
         Default::default(),
+        CancellationToken::new(),
     )
 }
 
@@ -126,35 +158,52 @@ pub struct FakeWfResponses {
     pub response_batches: Vec<ResponseType>,
 }
 
+// TODO: Should be all-internal to this module
 pub struct MocksHolder {
     client_bag: WorkerClientBag,
-    mock_worker: MockWorker,
-    // bidirectional mapping of run id / task token
-    pub outstanding_task_map: Option<Arc<RwLock<BiMap<String, TaskToken>>>>,
+    inputs: MockWorkerInputs,
+    pub outstanding_task_map: Option<OutstandingWFTMap>,
 }
 
 impl MocksHolder {
     pub fn worker_cfg(&mut self, mutator: impl FnOnce(&mut WorkerConfig)) {
-        mutator(&mut self.mock_worker.config);
+        mutator(&mut self.inputs.config);
+    }
+    pub fn set_act_poller(&mut self, poller: BoxedActPoller) {
+        self.inputs.act_poller = Some(poller);
+    }
+    /// Can be used for tests that need to avoid auto-shutdown due to running out of mock responses
+    pub 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 struct MockWorker {
-    pub wf_poller: BoxedWFPoller,
+pub struct MockWorkerInputs {
+    pub wft_stream: BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>,
     pub act_poller: Option<BoxedActPoller>,
     pub config: WorkerConfig,
 }
 
-impl Default for MockWorker {
+impl Default for MockWorkerInputs {
     fn default() -> Self {
-        Self::new(Box::from(mock_poller()))
+        Self::new_from_poller(Box::from(mock_poller()))
     }
 }
 
-impl MockWorker {
-    pub fn new(wf_poller: BoxedWFPoller) -> Self {
+impl MockWorkerInputs {
+    pub fn new(
+        wft_stream: BoxStream<'static, Result<ValidPollWFTQResponse, tonic::Status>>,
+    ) -> Self {
+        Self {
+            wft_stream,
+            act_poller: None,
+            config: test_worker_cfg().build().unwrap(),
+        }
+    }
+    pub fn new_from_poller(wf_poller: BoxedWFPoller) -> Self {
         Self {
-            wf_poller,
+            wft_stream: new_wft_poller(wf_poller, Default::default()).boxed(),
             act_poller: None,
             config: test_worker_cfg().build().unwrap(),
         }
@@ -162,54 +211,83 @@ impl MockWorker {
 }
 
 impl MocksHolder {
-    pub(crate) fn from_mock_worker(client_bag: WorkerClientBag, mock_worker: MockWorker) -> Self {
+    pub(crate) fn from_mock_worker(
+        client_bag: WorkerClientBag,
+        mock_worker: MockWorkerInputs,
+    ) -> Self {
         Self {
             client_bag,
-            mock_worker,
+            inputs: mock_worker,
             outstanding_task_map: None,
         }
     }
 
     /// Uses the provided list of tasks to create a mock poller for the `TEST_Q`
-    pub(crate) fn from_client_with_responses<WFT, ACT>(
+    pub(crate) fn from_client_with_activities<ACT>(
         client: impl WorkerClient + 'static,
-        wf_tasks: WFT,
         act_tasks: ACT,
     ) -> Self
     where
-        WFT: IntoIterator<Item = PollWorkflowTaskQueueResponse>,
-        ACT: IntoIterator<Item = PollActivityTaskQueueResponse>,
-        <WFT as IntoIterator>::IntoIter: Send + 'static,
+        ACT: IntoIterator<Item = QueueResponse<PollActivityTaskQueueResponse>>,
         <ACT as IntoIterator>::IntoIter: Send + 'static,
     {
-        let mock_poller = mock_poller_from_resps(wf_tasks);
+        let wft_stream = stream::pending().boxed();
         let mock_act_poller = mock_poller_from_resps(act_tasks);
-        let mock_worker = MockWorker {
-            wf_poller: mock_poller,
+        let mock_worker = MockWorkerInputs {
+            wft_stream,
             act_poller: Some(mock_act_poller),
             config: test_worker_cfg().build().unwrap(),
         };
         Self {
             client_bag: client.into(),
-            mock_worker,
+            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(
+        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 mock_worker = MockWorkerInputs {
+            wft_stream,
+            act_poller: None,
+            config: test_worker_cfg().build().unwrap(),
+        };
+        Self {
+            client_bag: client.into(),
+            inputs: mock_worker,
             outstanding_task_map: None,
         }
     }
 }
 
+// TODO: Un-pub ideally
 pub(crate) fn mock_poller_from_resps<T, I>(tasks: I) -> BoxedPoller<T>
 where
     T: Send + Sync + 'static,
-    I: IntoIterator<Item = T>,
+    I: IntoIterator<Item = QueueResponse<T>>,
     <I as IntoIterator>::IntoIter: Send + 'static,
 {
-    let mut mock_poller = mock_poller();
+    let mut mock_poller = mock_manual_poller();
     let mut tasks = tasks.into_iter();
     mock_poller.expect_poll().returning(move || {
         if let Some(t) = tasks.next() {
-            Some(Ok(t))
+            async move {
+                if let Some(f) = t.delay_until {
+                    f.await;
+                }
+                Some(Ok(t.resp))
+            }
+            .boxed()
         } else {
-            Some(Err(tonic::Status::cancelled(NO_MORE_WORK_ERROR_MSG)))
+            async { Some(Err(tonic::Status::cancelled(NO_MORE_WORK_ERROR_MSG))) }.boxed()
         }
     });
     Box::new(mock_poller) as BoxedPoller<T>
@@ -249,7 +327,7 @@ where
 pub(crate) fn build_multihist_mock_sg(
     hists: impl IntoIterator<Item = FakeWfResponses>,
     enforce_correct_number_of_polls: bool,
-    num_expected_fails: Option<usize>,
+    num_expected_fails: usize,
 ) -> MocksHolder {
     let mh = MockPollCfg::new(
         hists.into_iter().collect(),
@@ -275,11 +353,13 @@ pub(crate) fn single_hist_mock_sg(
 pub(crate) struct MockPollCfg {
     pub hists: Vec<FakeWfResponses>,
     pub enforce_correct_number_of_polls: bool,
-    pub num_expected_fails: Option<usize>,
+    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 expect_fail_wft_matcher:
         Box<dyn Fn(&TaskToken, &WorkflowTaskFailedCause, &Option<Failure>) -> bool + Send>,
+    pub completion_asserts: Option<Box<dyn Fn(&WorkflowTaskCompletion) + Send>>,
     /// 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.
@@ -290,14 +370,16 @@ impl MockPollCfg {
     pub fn new(
         hists: Vec<FakeWfResponses>,
         enforce_correct_number_of_polls: bool,
-        num_expected_fails: Option<usize>,
+        num_expected_fails: usize,
     ) -> Self {
         Self {
             hists,
             enforce_correct_number_of_polls,
             num_expected_fails,
+            num_expected_legacy_query_resps: 0,
             mock_client: mock_workflow_client(),
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
+            completion_asserts: None,
             using_rust_sdk: false,
         }
     }
@@ -314,19 +396,75 @@ impl MockPollCfg {
                 response_batches: resps.into_iter().map(Into::into).collect(),
             }],
             enforce_correct_number_of_polls: true,
-            num_expected_fails: None,
+            num_expected_fails: 0,
+            num_expected_legacy_query_resps: 0,
             mock_client,
             expect_fail_wft_matcher: Box::new(|_, _, _| true),
+            completion_asserts: None,
             using_rust_sdk: false,
         }
     }
 }
 
+#[derive(Default, Clone)]
+pub struct OutstandingWFTMap {
+    map: Arc<RwLock<BiMap<String, TaskToken>>>,
+    waker: Arc<Notify>,
+    all_work_delivered: Arc<AtomicBool>,
+}
+impl OutstandingWFTMap {
+    fn has_run(&self, run_id: &str) -> bool {
+        self.map.read().contains_left(run_id)
+    }
+    fn put_token(&self, run_id: String, token: TaskToken) {
+        self.map.write().insert(run_id, token);
+    }
+    fn release_token(&self, token: &TaskToken) {
+        self.map.write().remove_by_right(token);
+        self.waker.notify_one();
+    }
+    pub fn release_run(&self, run_id: &str) {
+        self.map.write().remove_by_left(run_id);
+        self.waker.notify_waiters();
+    }
+    pub fn all_work_delivered(&self) -> bool {
+        self.all_work_delivered.load(Ordering::Acquire)
+    }
+}
+
+struct EnsuresWorkDoneWFTStream {
+    inner: UnboundedReceiverStream<ValidPollWFTQResponse>,
+    all_work_was_completed: Arc<AtomicBool>,
+}
+impl Stream for EnsuresWorkDoneWFTStream {
+    type Item = ValidPollWFTQResponse;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.inner.poll_next_unpin(cx)
+    }
+}
+impl Drop for EnsuresWorkDoneWFTStream {
+    fn drop(&mut self) {
+        if !self.all_work_was_completed.load(Ordering::Acquire) && !std::thread::panicking() {
+            panic!("Not all workflow tasks were taken from mock!");
+        }
+    }
+}
+
 /// 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 {
     let mut task_q_resps: BTreeMap<String, VecDeque<_>> = BTreeMap::new();
-    let outstanding_wf_task_tokens = Arc::new(RwLock::new(BiMap::new()));
-    let mut correct_num_polls = None;
+    let all_work_delivered = if cfg.enforce_correct_number_of_polls && !cfg.using_rust_sdk {
+        Arc::new(AtomicBool::new(false))
+    } else {
+        Arc::new(AtomicBool::new(true))
+    };
+
+    let outstanding_wf_task_tokens = OutstandingWFTMap {
+        map: Arc::new(Default::default()),
+        waker: Arc::new(Default::default()),
+        all_work_delivered: all_work_delivered.clone(),
+    };
 
     for hist in cfg.hists {
         let full_hist_info = hist.hist.get_full_history_info().unwrap();
@@ -342,10 +480,6 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
             }
         }
 
-        if cfg.enforce_correct_number_of_polls && !cfg.using_rust_sdk {
-            *correct_num_polls.get_or_insert(0) += hist.response_batches.len();
-        }
-
         // Convert history batches into poll responses, while also tracking how many times a given
         // history has been returned so we can increment the associated attempt number on the WFT.
         // NOTE: This is hard to use properly with the `AfterEveryReply` testing eviction mode.
@@ -354,10 +488,10 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
         let mut attempts_at_task_num = HashMap::new();
         let responses: Vec<_> = hist
             .response_batches
-            .iter()
-            .map(|to_task_num| {
-                let cur_attempt = attempts_at_task_num.entry(to_task_num).or_insert(1);
-                let mut r = hist_to_poll_resp(&hist.hist, hist.wf_id.clone(), *to_task_num, TEST_Q);
+            .into_iter()
+            .map(|response| {
+                let cur_attempt = attempts_at_task_num.entry(response.hashable()).or_insert(1);
+                let mut r = hist_to_poll_resp(&hist.hist, hist.wf_id.clone(), response, TEST_Q);
                 r.attempt = *cur_attempt;
                 *cur_attempt += 1;
                 r
@@ -368,103 +502,168 @@ pub(crate) fn build_mock_pollers(mut cfg: MockPollCfg) -> MocksHolder {
         task_q_resps.insert(hist.wf_id, tasks);
     }
 
-    let mut mock_poller = mock_manual_poller();
     // The poller will return history from any workflow runs that do not have currently
     // outstanding tasks.
     let outstanding = outstanding_wf_task_tokens.clone();
-    let outstanding_wakeup_orig = Arc::new(Notify::new());
-    let outstanding_wakeup = outstanding_wakeup_orig.clone();
-    mock_poller
-        .expect_poll()
-        .times(correct_num_polls.map_or_else(|| RangeFull.into(), Into::<TimesRange>::into))
-        .returning(move || {
+    let outstanding_wakeup = outstanding.waker.clone();
+    let (wft_tx, wft_rx) = unbounded_channel();
+    tokio::task::spawn(async move {
+        loop {
             let mut resp = None;
-            for (_, tasks) in task_q_resps.iter_mut() {
+            let mut resp_iter = task_q_resps.iter_mut();
+            for (_, tasks) in &mut resp_iter {
                 // Must extract run id from a workflow task associated with this workflow
                 // TODO: Case where run id changes for same workflow id is not handled here
                 if let Some(t) = tasks.get(0) {
                     let rid = t.workflow_execution.as_ref().unwrap().run_id.clone();
-                    if !outstanding.read().contains_left(&rid) {
+                    if !outstanding.has_run(&rid) {
                         let t = tasks.pop_front().unwrap();
-                        outstanding
-                            .write()
-                            .insert(rid, TaskToken(t.task_token.clone()));
-                        resp = Some(Ok(t));
+                        outstanding.put_token(rid, TaskToken(t.task_token.clone()));
+                        resp = Some(t);
                         break;
                     }
                 }
             }
-            let outstanding_wakeup = outstanding_wakeup.clone();
-            async move {
-                if resp.is_some() {
-                    return resp;
-                }
+            let no_tasks_for_anyone = resp_iter.next().is_none();
 
-                if cfg.using_rust_sdk {
-                    // Simulate poll timeout, or just send an empty response and then try again
-                    // if we're told a new one might be ready.
-                    tokio::select! {
-                        _ = outstanding_wakeup.notified() => {}
-                        _ = tokio::time::sleep(Duration::from_secs(60)) => {}
-                    };
-                    Some(Ok(Default::default()))
-                } else {
-                    Some(Err(tonic::Status::cancelled(NO_MORE_WORK_ERROR_MSG)))
+            if let Some(resp) = resp {
+                if let Some(d) = resp.delay_until {
+                    d.await;
+                }
+                if wft_tx
+                    .send(
+                        resp.resp
+                            .try_into()
+                            .expect("Mock responses must be valid work"),
+                    )
+                    .is_err()
+                {
+                    dbg!("Exiting mock WFT task because rcv half of stream was dropped");
+                    break;
                 }
             }
-            .boxed()
-        });
-    let mock_worker = MockWorker::new(Box::from(mock_poller));
+
+            // No more work to do
+            if task_q_resps.values().all(|q| q.is_empty()) {
+                outstanding
+                    .all_work_delivered
+                    .store(true, Ordering::Release);
+                break;
+            }
+
+            if no_tasks_for_anyone {
+                tokio::select! {
+                    _ = outstanding_wakeup.notified() => {}
+                    _ = tokio::time::sleep(Duration::from_secs(60)) => {}
+                };
+            }
+        }
+    });
+    let mock_worker = MockWorkerInputs::new(
+        EnsuresWorkDoneWFTStream {
+            inner: UnboundedReceiverStream::new(wft_rx),
+            all_work_was_completed: all_work_delivered,
+        }
+        .map(Ok)
+        .boxed(),
+    );
 
     let outstanding = outstanding_wf_task_tokens.clone();
-    let outstanding_wakeup = outstanding_wakeup_orig.clone();
     cfg.mock_client
         .expect_complete_workflow_task()
         .returning(move |comp| {
-            outstanding.write().remove_by_right(&comp.task_token);
-            outstanding_wakeup.notify_one();
+            if let Some(ass) = cfg.completion_asserts.as_ref() {
+                // tee hee
+                ass(&comp)
+            }
+            outstanding.release_token(&comp.task_token);
             Ok(RespondWorkflowTaskCompletedResponse::default())
         });
     let outstanding = outstanding_wf_task_tokens.clone();
     cfg.mock_client
         .expect_fail_workflow_task()
         .withf(cfg.expect_fail_wft_matcher)
-        .times(
-            cfg.num_expected_fails
-                .map_or_else(|| RangeFull.into(), Into::<TimesRange>::into),
-        )
+        .times::<TimesRange>(cfg.num_expected_fails.into())
         .returning(move |tt, _, _| {
-            outstanding.write().remove_by_right(&tt);
-            outstanding_wakeup_orig.notify_one();
+            outstanding.release_token(&tt);
+            Ok(Default::default())
+        });
+    let outstanding = outstanding_wf_task_tokens.clone();
+    cfg.mock_client
+        .expect_respond_legacy_query()
+        .times::<TimesRange>(cfg.num_expected_legacy_query_resps.into())
+        .returning(move |tt, _| {
+            outstanding.release_token(&tt);
             Ok(Default::default())
         });
 
     MocksHolder {
         client_bag: cfg.mock_client.into(),
-        mock_worker,
+        inputs: mock_worker,
         outstanding_task_map: Some(outstanding_wf_task_tokens),
     }
 }
 
+pub struct QueueResponse<T> {
+    pub resp: T,
+    pub delay_until: Option<BoxFuture<'static, ()>>,
+}
+impl<T> From<T> for QueueResponse<T> {
+    fn from(resp: T) -> Self {
+        QueueResponse {
+            resp,
+            delay_until: None,
+        }
+    }
+}
+impl From<QueueResponse<PollWorkflowTaskQueueResponse>> for ResponseType {
+    fn from(qr: QueueResponse<PollWorkflowTaskQueueResponse>) -> Self {
+        ResponseType::Raw(qr.resp)
+    }
+}
+impl<T> Deref for QueueResponse<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.resp
+    }
+}
+impl<T> DerefMut for QueueResponse<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.resp
+    }
+}
+
 pub fn hist_to_poll_resp(
     t: &TestHistoryBuilder,
     wf_id: String,
     response_type: ResponseType,
     task_queue: impl Into<String>,
-) -> PollWorkflowTaskQueueResponse {
+) -> QueueResponse<PollWorkflowTaskQueueResponse> {
     let run_id = t.get_orig_run_id();
     let wf = WorkflowExecution {
         workflow_id: wf_id,
         run_id: run_id.to_string(),
     };
+    let mut delay_until = None;
     let hist_info = match response_type {
         ResponseType::ToTaskNum(tn) => t.get_history_info(tn).unwrap(),
         ResponseType::OneTask(tn) => t.get_one_wft(tn).unwrap(),
         ResponseType::AllHistory => t.get_full_history_info().unwrap(),
+        ResponseType::Raw(r) => {
+            return QueueResponse {
+                resp: r,
+                delay_until: None,
+            }
+        }
+        ResponseType::UntilResolved(fut, tn) => {
+            delay_until = Some(fut);
+            t.get_history_info(tn).unwrap()
+        }
     };
     let mut resp = hist_info.as_poll_wft_response(task_queue);
     resp.workflow_execution = Some(wf);
-    resp
+    QueueResponse { resp, delay_until }
 }
 
 type AsserterWithReply<'a> = (
@@ -472,6 +671,19 @@ type AsserterWithReply<'a> = (
     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.
 ///
@@ -489,7 +701,7 @@ pub(crate) async fn poll_and_reply<'a>(
 
 pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
     worker: &'a Worker,
-    outstanding_map: Option<Arc<RwLock<BiMap<String, TaskToken>>>>,
+    outstanding_map: Option<OutstandingWFTMap>,
     eviction_mode: WorkflowCachingPolicy,
     expect_and_reply: &'a [AsserterWithReply<'a>],
 ) {
@@ -515,6 +727,7 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
             let mut res = worker.poll_workflow_activation().await.unwrap();
             let contains_eviction = res.eviction_index();
 
+            let mut do_release = false;
             if let Some(eviction_job_ix) = contains_eviction {
                 // If the job list has an eviction, make sure it was the last item in the list
                 // then remove it, since in the tests we don't explicitly specify evict assertions
@@ -524,9 +737,7 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
                     "Eviction job was not last job in job list"
                 );
                 res.jobs.remove(eviction_job_ix);
-                if let Some(omap) = outstanding_map.as_ref() {
-                    omap.write().remove_by_left(&res.run_id);
-                }
+                do_release = true;
             }
 
             // TODO: Can remove this if?
@@ -549,6 +760,11 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
 
             worker.complete_workflow_activation(reply).await.unwrap();
 
+            if do_release {
+                if 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 contains_eviction.is_some() && !ends_execution {
@@ -560,7 +776,6 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
             }
 
             match eviction_mode {
-                WorkflowCachingPolicy::Sticky { .. } => unimplemented!(),
                 WorkflowCachingPolicy::NonSticky => (),
                 WorkflowCachingPolicy::AfterEveryReply => {
                     if evictions < expected_evictions {
@@ -575,7 +790,7 @@ pub(crate) async fn poll_and_reply_clears_outstanding_evicts<'a>(
     }
 
     assert_eq!(expected_fail_count, executed_failures.len());
-    assert_eq!(worker.outstanding_workflow_tasks(), 0);
+    assert_eq!(worker.outstanding_workflow_tasks().await, 0);
 }
 
 pub(crate) fn gen_assert_and_reply(
@@ -615,3 +830,20 @@ macro_rules! job_assert {
         }
     };
 }
+
+/// 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]
+macro_rules! advance_fut {
+    ($fut:ident) => {
+        ::futures::pin_mut!($fut);
+        {
+            let waker = ::futures::task::noop_waker();
+            let mut cx = core::task::Context::from_waker(&waker);
+            for _ in 0..10 {
+                assert_matches!($fut.poll_unpin(&mut cx), core::task::Poll::Pending);
+                ::tokio::task::yield_now().await;
+            }
+        }
+    };
+}