From 92049af24699fc2ba1b6df4fcbd63bd962e0a834 Mon Sep 17 00:00:00 2001
From: Lily Johnson <35852084+Lilyjjo@users.noreply.github.com>
Date: Mon, 1 Apr 2024 14:32:15 -0400
Subject: [PATCH] test(sequencer): add unit tests for src/accounts/state_ext
 (#871)

## Summary
Second set of state-ext unit tests. In total there are seven files which
need tests:
- `astria-sequencer/src/api_state_ext.rs`
- ~`astria-sequencer/src/state_ext.rs`~
- ~`astria-sequencer/src/accounts/state_ext.rs`~ (this PR)
- `astria-sequencer/src/asset/state_ext.rs`
- `astria-sequencer/src/authority/state_ext.rs`
- `astria-sequencer/src/bridge/state_ext.rs`
- `astria-sequencer/src/ibc/state_ext.rs`

This PR just tests the `src/accounts/state_ext.rs` file.

## Background
It is good to have unit tests to ensure that the database works as intended.

## Changes
Unit tests for the functionality in the file `src/accounts/state_ext.rs` were added.

## Testing
:)
---
 .../src/accounts/state_ext.rs                 | 480 ++++++++++++++++++
 1 file changed, 480 insertions(+)

diff --git a/crates/astria-sequencer/src/accounts/state_ext.rs b/crates/astria-sequencer/src/accounts/state_ext.rs
index 1898ec9a1..d340c0dcc 100644
--- a/crates/astria-sequencer/src/accounts/state_ext.rs
+++ b/crates/astria-sequencer/src/accounts/state_ext.rs
@@ -200,3 +200,483 @@ pub(crate) trait StateWriteExt: StateWrite {
 }
 
 impl<T: StateWrite> StateWriteExt for T {}
+
+#[cfg(test)]
+mod test {
+    use astria_core::sequencer::v1::{
+        account::AssetBalance,
+        asset::{
+            Denom,
+            Id,
+        },
+        Address,
+    };
+    use cnidarium::StateDelta;
+
+    use super::{
+        StateReadExt as _,
+        StateWriteExt as _,
+    };
+    use crate::asset;
+
+    #[tokio::test]
+    async fn get_account_nonce_uninitialized_returns_zero() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let nonce_expected = 0u32;
+
+        // uninitialized accounts return zero
+        assert_eq!(
+            state
+                .get_account_nonce(address)
+                .await
+                .expect("getting a non-initialized account's nonce should not fail"),
+            nonce_expected,
+            "returned nonce for non-initialized address was not zero"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_nonce_get_nonce_simple() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let nonce_expected = 0u32;
+
+        // can write new
+        state
+            .put_account_nonce(address, nonce_expected)
+            .expect("putting an account nonce should not fail");
+        assert_eq!(
+            state
+                .get_account_nonce(address)
+                .await
+                .expect("a nonce was written and must exist inside the database"),
+            nonce_expected,
+            "stored nonce was not what was expected"
+        );
+
+        // can rewrite with new value
+        let nonce_expected = 1u32;
+        state
+            .put_account_nonce(address, nonce_expected)
+            .expect("putting an account nonce should not fail");
+        assert_eq!(
+            state
+                .get_account_nonce(address)
+                .await
+                .expect("a new nonce was written and must exist inside the database"),
+            nonce_expected,
+            "updated nonce was not what was expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_nonce_get_nonce_complex() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let nonce_expected = 2u32;
+
+        // can write new
+        state
+            .put_account_nonce(address, nonce_expected)
+            .expect("putting an account nonce should not fail");
+        assert_eq!(
+            state
+                .get_account_nonce(address)
+                .await
+                .expect("a nonce was written and must exist inside the database"),
+            nonce_expected,
+            "stored nonce was not what was expected"
+        );
+
+        // writing additional account preserves first account's values
+        let address_1 = Address::try_from_slice(&[41u8; 20]).unwrap();
+        let nonce_expected_1 = 3u32;
+
+        state
+            .put_account_nonce(address_1, nonce_expected_1)
+            .expect("putting an account nonce should not fail");
+        assert_eq!(
+            state
+                .get_account_nonce(address_1)
+                .await
+                .expect("a new nonce was written and must exist inside the database"),
+            nonce_expected_1,
+            "additional account's nonce was not what was expected"
+        );
+        assert_eq!(
+            state
+                .get_account_nonce(address)
+                .await
+                .expect("a new nonce was written and must exist inside the database"),
+            nonce_expected,
+            "writing to a different account's nonce should not affect a different account's nonce"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_balance_uninitialized_returns_zero() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset = Id::from_denom("asset_0");
+        let amount_expected = 0u128;
+
+        // non-initialized accounts return zero
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting a non-initialized asset balance should not fail"),
+            amount_expected,
+            "returned balance for non-initialized asset balance was not zero"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_balance_simple() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset = Id::from_denom("asset_0");
+        let mut amount_expected = 1u128;
+
+        state
+            .put_account_balance(address, asset, amount_expected)
+            .expect("putting an account balance should not fail");
+
+        // can initialize
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected,
+            "returned balance for an asset balance did not match expected"
+        );
+
+        // can update balance
+        amount_expected = 2u128;
+
+        state
+            .put_account_balance(address, asset, amount_expected)
+            .expect("putting an asset balance for an account should not fail");
+
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected,
+            "returned balance for an asset did not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_balance_multiple_accounts() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset = Id::from_denom("asset_0");
+        let amount_expected = 1u128;
+
+        state
+            .put_account_balance(address, asset, amount_expected)
+            .expect("putting an account balance should not fail");
+
+        // able to write to account's storage
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected,
+            "returned balance for an asset did not match expected"
+        );
+
+        // writing to other accounts does not affect original account
+        // create needed variables
+        let address_1 = Address::try_from_slice(&[41u8; 20]).unwrap();
+        let amount_expected_1 = 2u128;
+
+        state
+            .put_account_balance(address_1, asset, amount_expected_1)
+            .expect("putting an account balance should not fail");
+        assert_eq!(
+            state
+                .get_account_balance(address_1, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected_1,
+            "returned balance for an asset did not match expected, changed during different \
+             account update"
+        );
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected,
+            "returned balance for an asset did not match expected, changed during different \
+             account update"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_balance_multiple_assets() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset_0 = Id::from_denom("asset_0");
+        let asset_1 = Id::from_denom("asset_1");
+        let amount_expected_0 = 1u128;
+        let amount_expected_1 = 2u128;
+
+        state
+            .put_account_balance(address, asset_0, amount_expected_0)
+            .expect("putting an account balance should not fail");
+        state
+            .put_account_balance(address, asset_1, amount_expected_1)
+            .expect("putting an account balance should not fail");
+
+        // wrote correct balances
+        assert_eq!(
+            state
+                .get_account_balance(address, asset_0)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected_0,
+            "returned balance for an asset did not match expected"
+        );
+        assert_eq!(
+            state
+                .get_account_balance(address, asset_1)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_expected_1,
+            "returned balance for an asset did not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn get_account_balances_uninitialized_ok() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+
+        // see that call was ok
+        let balances = state
+            .get_account_balances(address)
+            .await
+            .expect("retrieving account balances should not fail");
+        assert_eq!(balances, vec![]);
+    }
+
+    #[tokio::test]
+    async fn get_account_balances() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // need to set native asset in order to use `get_account_balances()`
+        crate::asset::initialize_native_asset("asset_0");
+
+        let asset_0 = Id::from_denom("asset_0");
+        let asset_1 = Id::from_denom("asset_1");
+        let asset_2 = Id::from_denom("asset_2");
+
+        // also need to add assets to the ibc state
+        asset::state_ext::StateWriteExt::put_ibc_asset(
+            &mut state,
+            asset_0,
+            &Denom::from_base_denom("asset_0"),
+        )
+        .expect("should be able to call other trait method on state object");
+        asset::state_ext::StateWriteExt::put_ibc_asset(
+            &mut state,
+            asset_1,
+            &Denom::from_base_denom("asset_1"),
+        )
+        .expect("should be able to call other trait method on state object");
+        asset::state_ext::StateWriteExt::put_ibc_asset(
+            &mut state,
+            asset_2,
+            &Denom::from_base_denom("asset_2"),
+        )
+        .expect("should be able to call other trait method on state object");
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let amount_expected_0 = 1u128;
+        let amount_expected_1 = 2u128;
+        let amount_expected_2 = 3u128;
+
+        // add balances to the account
+        state
+            .put_account_balance(address, asset_0, amount_expected_0)
+            .expect("putting an account balance should not fail");
+        state
+            .put_account_balance(address, asset_1, amount_expected_1)
+            .expect("putting an account balance should not fail");
+        state
+            .put_account_balance(address, asset_2, amount_expected_2)
+            .expect("putting an account balance should not fail");
+
+        let mut balances = state
+            .get_account_balances(address)
+            .await
+            .expect("retrieving account balances should not fail");
+        balances.sort_by(|a, b| a.balance.cmp(&b.balance));
+        assert_eq!(
+            balances,
+            vec![
+                AssetBalance {
+                    denom: Denom::from_base_denom("asset_0"),
+                    balance: amount_expected_0,
+                },
+                AssetBalance {
+                    denom: Denom::from_base_denom("asset_1"),
+                    balance: amount_expected_1,
+                },
+                AssetBalance {
+                    denom: Denom::from_base_denom("asset_2"),
+                    balance: amount_expected_2,
+                },
+            ]
+        );
+    }
+
+    #[tokio::test]
+    async fn increase_balance_from_uninitialized() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset = Id::from_denom("asset_0");
+        let amount_increase = 2u128;
+
+        state
+            .increase_balance(address, asset, amount_increase)
+            .await
+            .expect("increasing account balance for uninitialized account should be ok");
+
+        // correct balance was set
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_increase,
+            "returned balance for an asset balance did not match expected"
+        );
+
+        state
+            .increase_balance(address, asset, amount_increase)
+            .await
+            .expect("increasing account balance for initialized account should be ok");
+
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_increase * 2,
+            "returned balance for an asset balance did not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn decrease_balance_enough_funds() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset = Id::from_denom("asset_0");
+        let amount_increase = 2u128;
+
+        state
+            .increase_balance(address, asset, amount_increase)
+            .await
+            .expect("increasing account balance for uninitialized account should be ok");
+
+        // correct balance was set
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            amount_increase,
+            "returned balance for an asset balance did not match expected"
+        );
+
+        // decrease balance
+        state
+            .decrease_balance(address, asset, amount_increase)
+            .await
+            .expect("decreasing account balance for initialized account should be ok");
+
+        assert_eq!(
+            state
+                .get_account_balance(address, asset)
+                .await
+                .expect("getting an asset balance should not fail"),
+            0,
+            "returned balance for an asset balance did not match expected"
+        );
+    }
+
+    #[tokio::test]
+    async fn decrease_balance_not_enough_funds() {
+        let storage = cnidarium::TempStorage::new().await.unwrap();
+        let snapshot = storage.latest_snapshot();
+        let mut state = StateDelta::new(snapshot);
+
+        // create needed variables
+        let address = Address::try_from_slice(&[42u8; 20]).unwrap();
+        let asset = Id::from_denom("asset_0");
+        let amount_increase = 2u128;
+
+        // give initial balance
+        state
+            .increase_balance(address, asset, amount_increase)
+            .await
+            .expect("increasing account balance for uninitialized account should be ok");
+
+        // decrease balance
+        state
+            .decrease_balance(address, asset, amount_increase + 1)
+            .await
+            .expect_err("should not be able to subtract larger balance than what existed");
+    }
+}