@temporalio/core-bridge 0.19.2 → 0.20.2

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

@@ -1,15 +1,14 @@
 use crate::{
     test_help::{
-        canned_histories, hist_to_poll_resp, mock_core, MocksHolder, ResponseType, TEST_Q,
+        canned_histories, hist_to_poll_resp, mock_worker, MocksHolder, ResponseType, TEST_Q,
     },
-    Core,
+    worker::client::mocks::mock_workflow_client,
 };
 use std::{
     collections::{HashMap, VecDeque},
     time::Duration,
 };
-use temporal_client::mocks::mock_gateway;
-
+use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::{
     coresdk::{
         workflow_activation::{workflow_activation_job, WorkflowActivationJob},
@@ -17,6 +16,7 @@ use temporal_sdk_core_protos::{
         workflow_completion::WorkflowActivationCompletion,
     },
     temporal::api::{
+        common::v1::Payload,
         failure::v1::Failure,
         history::v1::History,
         query::v1::WorkflowQuery,
@@ -35,6 +35,8 @@ async fn legacy_query(#[case] include_history: bool) {
     let wfid = "fake_wf_id";
     let query_resp = "response";
     let t = canned_histories::single_timer("1");
+    let mut header = HashMap::new();
+    header.insert("head".to_string(), Payload::from(b"er"));
     let tasks = VecDeque::from(vec![
         hist_to_poll_resp(&t, wfid.to_owned(), 1.into(), TEST_Q.to_string()),
         {
@@ -42,7 +44,7 @@ async fn legacy_query(#[case] include_history: bool) {
             pr.query = Some(WorkflowQuery {
                 query_type: "query-type".to_string(),
                 query_args: Some(b"hi".into()),
-                header: None,
+                header: Some(header.into()),
             });
             if !include_history {
                 pr.history = Some(History { events: vec![] });
@@ -51,34 +53,35 @@ async fn legacy_query(#[case] include_history: bool) {
         },
         hist_to_poll_resp(&t, wfid.to_owned(), 2.into(), TEST_Q.to_string()),
     ]);
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
-    mock_gateway
+    mock_client
         .expect_respond_legacy_query()
         .times(1)
         .returning(move |_, _| Ok(RespondQueryTaskCompletedResponse::default()));
 
-    let mut mock = MocksHolder::from_gateway_with_responses(mock_gateway, tasks, vec![]);
+    let mut mock = MocksHolder::from_client_with_responses(mock_client, tasks, vec![]);
     if !include_history {
-        mock.worker_cfg(TEST_Q, |wc| wc.max_cached_workflows = 10);
+        mock.worker_cfg(|wc| wc.max_cached_workflows = 10);
     }
-    let core = mock_core(mock);
+    let worker = mock_worker(mock);
 
     let first_wft = || async {
-        let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
-        core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-            TEST_Q,
-            task.run_id,
-            start_timer_cmd(1, Duration::from_secs(1)),
-        ))
-        .await
-        .unwrap();
+        let task = worker.poll_workflow_activation().await.unwrap();
+        worker
+            .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+                task.run_id,
+                start_timer_cmd(1, Duration::from_secs(1)),
+            ))
+            .await
+            .unwrap();
     };
     let clear_eviction = || async {
-        let t = core.poll_workflow_activation(TEST_Q).await.unwrap();
-        core.complete_workflow_activation(WorkflowActivationCompletion::empty(TEST_Q, t.run_id))
+        let t = worker.poll_workflow_activation().await.unwrap();
+        worker
+            .complete_workflow_activation(WorkflowActivationCompletion::empty(t.run_id))
             .await
             .unwrap();
     };
@@ -90,7 +93,7 @@ async fn legacy_query(#[case] include_history: bool) {
         first_wft().await;
     }
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = worker.poll_workflow_activation().await.unwrap();
     // Poll again, and we end up getting a `query` field query response
     let query = assert_matches!(
         task.jobs.as_slice(),
@@ -98,44 +101,45 @@ async fn legacy_query(#[case] include_history: bool) {
             variant: Some(workflow_activation_job::Variant::QueryWorkflow(q)),
         }] => q
     );
+    assert_eq!(query.headers.get("head").unwrap(), &b"er".into());
     // Complete the query
-    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-        TEST_Q,
-        task.run_id,
-        QueryResult {
-            query_id: query.query_id.clone(),
-            variant: Some(
-                QuerySuccess {
-                    response: Some(query_resp.into()),
-                }
-                .into(),
-            ),
-        }
-        .into(),
-    ))
-    .await
-    .unwrap();
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+            task.run_id,
+            QueryResult {
+                query_id: query.query_id.clone(),
+                variant: Some(
+                    QuerySuccess {
+                        response: Some(query_resp.into()),
+                    }
+                    .into(),
+                ),
+            }
+            .into(),
+        ))
+        .await
+        .unwrap();
 
     if include_history {
         clear_eviction().await;
         first_wft().await;
     }
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = worker.poll_workflow_activation().await.unwrap();
     assert_matches!(
         task.jobs.as_slice(),
         [WorkflowActivationJob {
             variant: Some(workflow_activation_job::Variant::FireTimer(_)),
         }]
     );
-    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
-        TEST_Q,
-        task.run_id,
-        vec![CompleteWorkflowExecution { result: None }.into()],
-    ))
-    .await
-    .unwrap();
-    core.shutdown().await;
+    worker
+        .complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+            task.run_id,
+            vec![CompleteWorkflowExecution { result: None }.into()],
+        ))
+        .await
+        .unwrap();
+    worker.shutdown().await;
 }
 
 #[rstest::rstest]
@@ -146,6 +150,8 @@ async fn new_queries(#[case] num_queries: usize) {
     let wfid = "fake_wf_id";
     let query_resp = "response";
     let t = canned_histories::single_timer("1");
+    let mut header = HashMap::new();
+    header.insert("head".to_string(), Payload::from(b"er"));
     let tasks = VecDeque::from(vec![
         hist_to_poll_resp(&t, wfid.to_owned(), 1.into(), TEST_Q.to_string()),
         {
@@ -157,33 +163,32 @@ async fn new_queries(#[case] num_queries: usize) {
                     WorkflowQuery {
                         query_type: "query-type".to_string(),
                         query_args: Some(b"hi".into()),
-                        header: None,
+                        header: Some(header.clone().into()),
                     },
                 );
             }
             pr
         },
     ]);
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
-    mock_gateway.expect_respond_legacy_query().times(0);
+    mock_client.expect_respond_legacy_query().times(0);
 
-    let mut mock = MocksHolder::from_gateway_with_responses(mock_gateway, tasks, vec![]);
-    mock.worker_cfg(TEST_Q, |wc| wc.max_cached_workflows = 10);
-    let core = mock_core(mock);
+    let mut mock = MocksHolder::from_client_with_responses(mock_client, tasks, vec![]);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 10);
+    let core = mock_worker(mock);
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-        TEST_Q,
         task.run_id,
         start_timer_cmd(1, Duration::from_secs(1)),
     ))
     .await
     .unwrap();
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     assert_matches!(
         task.jobs[0],
         WorkflowActivationJob {
@@ -194,7 +199,9 @@ async fn new_queries(#[case] num_queries: usize) {
         assert_matches!(
             task.jobs[i],
             WorkflowActivationJob {
-                variant: Some(workflow_activation_job::Variant::QueryWorkflow(_)),
+                variant: Some(workflow_activation_job::Variant::QueryWorkflow(ref q)),
+            } => {
+                assert_eq!(q.headers.get("head").unwrap(), &b"er".into());
             }
         );
     }
@@ -214,7 +221,6 @@ async fn new_queries(#[case] num_queries: usize) {
         .collect();
     qresults.push(CompleteWorkflowExecution { result: None }.into());
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
-        TEST_Q,
         task.run_id,
         qresults,
     ))
@@ -241,29 +247,28 @@ async fn legacy_query_failure_on_wft_failure() {
         },
         hist_to_poll_resp(&t, wfid.to_owned(), 2.into(), TEST_Q.to_string()),
     ]);
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
-    mock_gateway
+    mock_client
         .expect_respond_legacy_query()
         .times(1)
         .returning(move |_, _| Ok(RespondQueryTaskCompletedResponse::default()));
 
-    let mut mock = MocksHolder::from_gateway_with_responses(mock_gateway, tasks, vec![]);
-    mock.worker_cfg(TEST_Q, |wc| wc.max_cached_workflows = 10);
-    let core = mock_core(mock);
+    let mut mock = MocksHolder::from_client_with_responses(mock_client, tasks, vec![]);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 10);
+    let core = mock_worker(mock);
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-        TEST_Q,
         task.run_id,
         start_timer_cmd(1, Duration::from_secs(1)),
     ))
     .await
     .unwrap();
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     // Poll again, and we end up getting a `query` field query response
     assert_matches!(
         task.jobs.as_slice(),
@@ -273,7 +278,6 @@ async fn legacy_query_failure_on_wft_failure() {
     );
     // Fail wft which should result in query being failed
     core.complete_workflow_activation(WorkflowActivationCompletion::fail(
-        TEST_Q,
         task.run_id,
         Failure {
             message: "Ahh i broke".to_string(),
@@ -321,30 +325,28 @@ async fn legacy_query_after_complete(#[values(false, true)] full_history: bool)
         query_with_hist_task.clone(),
         query_with_hist_task,
     ]);
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_complete_workflow_task()
         .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
-    mock_gateway
+    mock_client
         .expect_respond_legacy_query()
         .times(2)
         .returning(move |_, _| Ok(RespondQueryTaskCompletedResponse::default()));
 
-    let mut mock = MocksHolder::from_gateway_with_responses(mock_gateway, tasks, vec![]);
-    mock.worker_cfg(TEST_Q, |wc| wc.max_cached_workflows = 10);
-    let core = mock_core(mock);
+    let mut mock = MocksHolder::from_client_with_responses(mock_client, tasks, vec![]);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 10);
+    let core = mock_worker(mock);
 
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-        TEST_Q,
         task.run_id,
         start_timer_cmd(1, Duration::from_secs(1)),
     ))
     .await
     .unwrap();
-    let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
     core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
-        TEST_Q,
         task.run_id,
         vec![CompleteWorkflowExecution { result: None }.into()],
     ))
@@ -353,7 +355,7 @@ async fn legacy_query_after_complete(#[values(false, true)] full_history: bool)
 
     // We should get queries two times
     for _ in 1..=2 {
-        let task = core.poll_workflow_activation(TEST_Q).await.unwrap();
+        let task = core.poll_workflow_activation().await.unwrap();
         let query = assert_matches!(
             task.jobs.as_slice(),
             [WorkflowActivationJob {
@@ -361,7 +363,6 @@ async fn legacy_query_after_complete(#[values(false, true)] full_history: bool)
             }] => q
         );
         core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-            TEST_Q,
             task.run_id,
             QueryResult {
                 query_id: query.query_id.clone(),