test(sequencer): add unit tests for src/api_state_ext

crates/astria-sequencer/src/api_state_ext.rs

@@ -375,3 +375,350 @@ pub(crate) trait StateWriteExt: StateWrite {
 }
 
 impl<T: StateWrite> StateWriteExt for T {}
+
+#[cfg(test)]
+mod test {
+
+    use std::collections::HashMap;
+
+    use astria_core::{
+        sequencer::v1::{
+            asset::Id,
+            block::{
+                Deposit,
+                SequencerBlock,
+            },
+            Address,
+            RollupId,
+        },
+        Protobuf,
+    };
+    use cnidarium::StateDelta;
+    use rand::Rng;
+    use sha2::{
+        Digest as _,
+        Sha256,
+    };
+    use tendermint::{
+        account,
+        block::{
+            header::Version,
+            Header,
+        },
+        AppHash,
+        Hash,
+        Time,
+    };
+
+    use super::*;
+    use crate::proposal::commitment::generate_rollup_datas_commitment;
+
+    // creates new sequencer block, optionally shifting all values except the height by 1
+    fn make_test_sequencer_block(height: u32) -> SequencerBlock {
+        let mut rng = rand::thread_rng();
+
+        // create cometbft header
+        let mut header = Header {
+            app_hash: AppHash::default(),
+            chain_id: "test".to_string().try_into().unwrap(),
+            consensus_hash: Hash::default(),
+            data_hash: Some(Hash::default()),
+            evidence_hash: None,
+            height: height.into(),
+            last_block_id: None,
+            last_commit_hash: None,
+            last_results_hash: None,
+            next_validators_hash: Hash::default(),
+            proposer_address: account::Id::try_from([0u8; 20].to_vec()).unwrap(),
+            time: Time::now(),
+            validators_hash: Hash::default(),
+            version: Version {
+                app: 0,
+                block: 0,
+            },
+        };
+
+        // create inner rollup id/tx data
+        let mut deposits = HashMap::new();
+        for _ in 0..2 {
+            let rollup_id = RollupId::new(rng.gen());
+            let bridge_address = Address::try_from_slice(&[rng.gen(); 20]).unwrap();
+            let amount = rng.gen::<u128>();
+            let asset_id = Id::from_denom(&rng.gen::<u8>().to_string());
+            let destination_chain_address = rng.gen::<u8>().to_string();
+            let deposit = Deposit::new(
+                bridge_address,
+                rollup_id,
+                amount,
+                asset_id,
+                destination_chain_address,
+            );
+            deposits.insert(rollup_id, vec![deposit]);
+        }
+
+        // create block's other data
+        let commitments = generate_rollup_datas_commitment(&[], deposits.clone());
+        let block_data = vec![
+            commitments.rollup_datas_root.to_vec(),
+            commitments.rollup_ids_root.to_vec(),
+        ];
+        let data_hash = merkle::Tree::from_leaves(block_data.iter().map(Sha256::digest)).root();
+        header.data_hash = Some(Hash::try_from(data_hash.to_vec()).unwrap());
+        header.height = height.into();
+        SequencerBlock::try_from_cometbft_header_and_data(header, block_data, deposits).unwrap()
+    }
+
+    #[tokio::test]
+    async fn put_sequencer_block() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // can write one
+        let block_0 = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block_0.clone())
+            .expect("writing block to database should work");
+
+        assert_eq!(
+            state
+                .get_sequencer_block_by_height(block_0.height().into())
+                .await
+                .expect("a block was written to the database and should exist"),
+            block_0,
+            "stored block does not match expected"
+        );
+
+        // can write another and both are ok
+        let block_1 = make_test_sequencer_block(3u32);
+        state
+            .put_sequencer_block(block_1.clone())
+            .expect("writing another block to database should work");
+        assert_eq!(
+            state
+                .get_sequencer_block_by_height(block_0.height().into())
+                .await
+                .expect("a block was written to the database and should exist"),
+            block_0,
+            "original stored block does not match expected"
+        );
+        assert_eq!(
+            state
+                .get_sequencer_block_by_height(block_1.height().into())
+                .await
+                .expect("a block was written to the database and should exist"),
+            block_1,
+            "additionally stored block does not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn put_sequencer_block_update() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write original block
+        let mut block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+        assert_eq!(
+            state
+                .get_sequencer_block_by_height(block.height().into())
+                .await
+                .expect("a block was written to the database and should exist"),
+            block,
+            "stored block does not match expected"
+        );
+
+        // write to same height but with new values
+        block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block update to database should work");
+
+        // block was updates
+        assert_eq!(
+            state
+                .get_sequencer_block_by_height(block.height().into())
+                .await
+                .expect("a block was written to the database and should exist"),
+            block,
+            "updated stored block does not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_block_hash_by_height() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write block
+        let block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+
+        // grab block hash by block height
+        assert_eq!(
+            state
+                .get_block_hash_by_height(block.height().into())
+                .await
+                .expect(
+                    "a block was written to the database and we should be able to query its block \
+                     hash by height"
+                ),
+            block.block_hash(),
+            "stored block hash does not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_sequencer_block_header_by_hash() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write block
+        let block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+
+        // grab block header by block hash
+        assert_eq!(
+            state
+                .get_sequencer_block_header_by_hash(block.block_hash().as_ref())
+                .await
+                .expect(
+                    "a block was written to the database and we should be able to query its block \
+                     header by block hash"
+                ),
+            block.header().clone(),
+            "stored block header does not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_rollup_ids_by_block_hash() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write block
+        let block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+
+        // grab rollup ids by block hash
+        let stored_rollup_ids = state
+            .get_rollup_ids_by_block_hash(block.block_hash().as_ref())
+            .await
+            .expect(
+                "a block was written to the database and we should be able to query its rollup ids",
+            );
+        let original_rollup_ids: Vec<RollupId> =
+            block.rollup_transactions().keys().copied().collect();
+        assert_eq!(
+            stored_rollup_ids, original_rollup_ids,
+            "stored rollup ids do not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_sequencer_block_by_hash() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write block
+        let block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+
+        // grab block by block hash
+        assert_eq!(
+            state
+                .get_sequencer_block_by_hash(block.block_hash().as_ref())
+                .await
+                .expect(
+                    "a block was written to the database and we should be able to query its block \
+                     by block hash"
+                ),
+            block,
+            "stored block does not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_rollup_data() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write block
+        let block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+
+        // get written rollup id and data
+        let rollup_id = block
+            .rollup_transactions()
+            .keys()
+            .copied()
+            .collect::<Vec<RollupId>>()[0];
+        let rollup_data = block.rollup_transactions().get(&rollup_id).unwrap();
+
+        // grab rollup's data by block hash
+        let stored_rollup_data = state
+            .get_rollup_data(block.block_hash().as_ref(), &rollup_id)
+            .await
+            .expect(
+                "a block was written to the database and we should be able to query the data for \
+                 a rollup",
+            );
+        assert_eq!(
+            stored_rollup_data, *rollup_data,
+            "stored rollup data does not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_block_proofs_by_block_hash() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // write block
+        let block = make_test_sequencer_block(2u32);
+        state
+            .put_sequencer_block(block.clone())
+            .expect("writing block to database should work");
+
+        // get written proofs
+        let transactions_proof = block
+            .clone()
+            .into_parts()
+            .rollup_transactions_proof
+            .into_raw();
+        let ids_proof = block.clone().into_parts().rollup_ids_proof.into_raw();
+
+        // grab rollup's stored proofs
+        let stored_proofs = state
+            .get_block_proofs_by_block_hash(block.block_hash().as_ref())
+            .await
+            .expect(
+                "a block was written to the database and we should be able to query its proof data",
+            );
+        assert_eq!(
+            (transactions_proof, ids_proof),
+            stored_proofs,
+            "stored proofs do not match expected"
+        );
+    }
+}