runtime-sdk/testing: allow mock with custom keys

oasisprotocol · Oct 18, 2024 · ed33478 · ed33478
1 parent 73bd772
commit ed33478
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/runtime-sdk/src/testing/keymanager.rs b/runtime-sdk/src/testing/keymanager.rs
@@ -15,8 +15,8 @@ use crate::{
 
 #[derive(Default)]
 pub struct MockKeyManagerClient {
-    keys: Mutex<HashMap<KeyPairId, KeyPair>>,
-    ephemeral_keys: Mutex<HashMap<KeyPairId, KeyPair>>,
+    pub keys: Mutex<HashMap<KeyPairId, KeyPair>>,
+    pub ephemeral_keys: Mutex<HashMap<KeyPairId, KeyPair>>,
 }
 
 impl Clone for MockKeyManagerClient {

diff --git a/runtime-sdk/src/testing/mock.rs b/runtime-sdk/src/testing/mock.rs
@@ -101,6 +101,20 @@ impl Mock {
         )
     }
 
+    /// Create a new mock dispatch context.
+    pub fn create_ctx_with_key_manager(&mut self, key_manager: Box<dyn KeyManager>) -> RuntimeBatchContext<'_, EmptyRuntime> {
+        RuntimeBatchContext::new(
+            &self.host_info,
+            Some(key_manager),
+            &self.runtime_header,
+            &self.runtime_round_results,
+            &self.consensus_state,
+            &self.history,
+            self.epoch,
+            self.max_messages,
+        )
+    }
+
     /// Create an instance with the given local configuration.
     pub fn with_local_config(local_config: BTreeMap<String, cbor::Value>) -> Self {
         // Ensure a current state is always available during tests. Note that one can always use a