From 880f566e9cd9022452107cec454f20b62450bace Mon Sep 17 00:00:00 2001
From: Tomasz Polaczyk <tmpolaczyk@gmail.com>
Date: Thu, 22 Aug 2024 13:25:53 +0200
Subject: [PATCH] Add tests

---
 .../authorities-noting/src/mock.rs            | 17 ++++++++-
 .../authorities-noting/src/tests.rs           | 19 ++++++++++
 test-sproof-builder/src/lib.rs                | 37 +++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/container-chain-pallets/authorities-noting/src/mock.rs b/container-chain-pallets/authorities-noting/src/mock.rs
index 2106a0f..9a69aae 100644
--- a/container-chain-pallets/authorities-noting/src/mock.rs
+++ b/container-chain-pallets/authorities-noting/src/mock.rs
@@ -180,6 +180,7 @@ pub struct BlockTests {
     relay_sproof_builder_hook:
         Option<Box<dyn Fn(&BlockTests, RelayChainBlockNumber, &mut ParaHeaderSproofBuilder)>>,
     orchestrator_storage_proof: Option<StorageProof>,
+    relay_storage_proof: Option<(H256, StorageProof)>,
     skip_inherent_insertion: bool,
 }
 
@@ -218,6 +219,12 @@ where {
         self
     }
 
+    pub fn with_relay_storage_proof(mut self, root: H256, proof: StorageProof) -> Self
+where {
+        self.relay_storage_proof = Some((root, proof));
+        self
+    }
+
     pub fn skip_inherent_insertion(mut self) -> Self {
         self.skip_inherent_insertion = true;
         self
@@ -242,9 +249,17 @@ where {
                     hook(self, *n as RelayChainBlockNumber, &mut sproof_builder);
                 }
 
-                let (relay_parent_storage_root, relay_chain_state) =
+                let (mut relay_parent_storage_root, mut relay_chain_state) =
                     sproof_builder.into_state_root_and_proof();
 
+                if let Some((root, proof)) = self.relay_storage_proof.clone() {
+                    if self.relay_sproof_builder_hook.is_some() {
+                        panic!("Cannot set both relay_sproof_builder_hook and relay_storage_proof, remove hook as it will be ignored anyway")
+                    }
+                    relay_parent_storage_root = root;
+                    relay_chain_state = proof;
+                }
+
                 // We write relay storage root in mock storage.
                 frame_support::storage::unhashed::put(
                     MOCK_RELAY_ROOT_KEY,
diff --git a/container-chain-pallets/authorities-noting/src/tests.rs b/container-chain-pallets/authorities-noting/src/tests.rs
index 4265365..3e86af9 100644
--- a/container-chain-pallets/authorities-noting/src/tests.rs
+++ b/container-chain-pallets/authorities-noting/src/tests.rs
@@ -114,6 +114,25 @@ fn test_authorities_insertion_wrong_para_id() {
         });
 }
 
+#[test]
+fn test_authorities_insertion_right_para_id_solochain() {
+    let mut assignment = AuthorityAssignmentSproofBuilder::<u64>::default();
+    assignment
+        .authority_assignment
+        .container_chains
+        .insert(ParachainId::get(), vec![10u64, 11u64]);
+
+    let (orchestrator_chain_root, orchestrator_chain_state) =
+        assignment.into_state_root_and_proof_solochain();
+
+    BlockTests::new()
+        .with_relay_storage_proof(orchestrator_chain_root, orchestrator_chain_state)
+        .with_orchestrator_storage_proof(sp_trie::StorageProof::empty())
+        .add(1, || {
+            assert_eq!(AuthoritiesNoting::authorities(), vec![10u64, 11u64]);
+        });
+}
+
 #[test]
 #[should_panic(
     expected = "Orchestrator chain authorities data needs to be present in every block!"
diff --git a/test-sproof-builder/src/lib.rs b/test-sproof-builder/src/lib.rs
index e624685..0ce7680 100644
--- a/test-sproof-builder/src/lib.rs
+++ b/test-sproof-builder/src/lib.rs
@@ -209,4 +209,41 @@ impl<T: Encode> AuthorityAssignmentSproofBuilder<T> {
 
         (root, proof)
     }
+
+    pub fn into_state_root_and_proof_solochain(
+        self,
+    ) -> (
+        cumulus_primitives_core::relay_chain::Hash,
+        sp_state_machine::StorageProof,
+    ) {
+        let (db, root) =
+            PrefixedMemoryDB::<HashingFor<cumulus_primitives_core::relay_chain::Block>>::default_with_root();
+
+        let state_version = Default::default();
+        let mut backend = sp_state_machine::TrieBackendBuilder::new(db, root).build();
+        let mut relevant_keys = Vec::new();
+
+        let mut insert = |key: Vec<u8>, value: Vec<u8>| {
+            relevant_keys.push(key.clone());
+            backend.insert(vec![(None, vec![(key, Some(value))])], state_version);
+        };
+
+        insert(
+            well_known_keys::SESSION_INDEX.to_vec(),
+            self.session_index.encode(),
+        );
+        insert(
+            well_known_keys::authority_assignment_for_session(
+                self.session_index,
+                Some(well_known_keys::SOLOCHAIN_AUTHORITY_ASSIGNMENT_PREFIX),
+            )
+            .to_vec(),
+            self.authority_assignment.encode(),
+        );
+
+        let root = *backend.root();
+        let proof = sp_state_machine::prove_read(backend, relevant_keys).expect("prove read");
+
+        (root, proof)
+    }
 }