TestKit: add executeQuery sessionAuth config

robsdedude · Apr 13, 2024 · e3061f9 · e3061f9
1 parent dbc0c59
commit e3061f9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/testkit_backend/src/testkit_backend/driver_holder.rs b/testkit_backend/src/testkit_backend/driver_holder.rs
@@ -655,6 +655,7 @@ impl DriverHolderRunner {
                     bookmark_manager,
                     tx_meta,
                     timeout,
+                    auth,
                 }) => {
                     let mut builder = self.driver.execute_query(query);
                     if let Some(params) = params {
@@ -684,6 +685,9 @@ impl DriverHolderRunner {
                     if let Some(timeout) = timeout {
                         builder = builder.with_transaction_timeout(timeout);
                     }
+                    if let Some(auth) = auth {
+                        builder = builder.with_session_auth(auth);
+                    }
                     let result = builder
                         .run_with_retry(ExponentialBackoff::default())
                         .map_err(Into::into);
@@ -1051,6 +1055,7 @@ pub(super) struct ExecuteQuery {
     pub(crate) bookmark_manager: ExecuteQueryBookmarkManager,
     pub(super) tx_meta: Option<HashMap<String, ValueSend>>,
     pub(super) timeout: Option<TransactionTimeout>,
+    pub(super) auth: Option<Arc<AuthToken>>,
 }
 
 impl From<ExecuteQuery> for Command {

diff --git a/testkit_backend/src/testkit_backend/requests.rs b/testkit_backend/src/testkit_backend/requests.rs
@@ -599,6 +599,8 @@ pub(super) struct ExecuteQueryConfig {
     bookmark_manager_id: Option<ExecuteQueryBmmId>,
     tx_meta: Option<HashMap<String, CypherValue>>,
     timeout: Option<i64>,
+    #[serde(rename = "authorizationToken")]
+    auth: MaybeTestKitAuth,
 }
 
 #[derive(Deserialize, Debug)]
@@ -1466,6 +1468,7 @@ impl Request {
             bookmark_manager_id,
             tx_meta,
             timeout,
+            auth,
         } = config;
         let data = backend.data.borrow_mut();
         let driver_holder = get_driver(&data, &driver_id)?;
@@ -1483,6 +1486,7 @@ impl Request {
             .map(cypher_value_map_to_value_send_map)
             .transpose()?;
         let timeout = read_transaction_timeout(timeout)?;
+        let auth = auth.0.map(|auth| Arc::new(auth.into()));
         let result = driver_holder
             .execute_query(ExecuteQuery {
                 query,
@@ -1493,6 +1497,7 @@ impl Request {
                 bookmark_manager,
                 tx_meta,
                 timeout,
+                auth,
             })
             .result?;
         let response: Response = result.try_into()?;