From d2e2788cad897046599efe2624dc4a05b4f10341 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Mon, 5 Sep 2022 12:20:34 +0200
Subject: [PATCH] runtime: Add beacon::MutableState for easier testing

---
 runtime/src/consensus/state/beacon.rs | 81 ++++++++++++++++++++++++++-
 1 file changed, 80 insertions(+), 1 deletion(-)

diff --git a/runtime/src/consensus/state/beacon.rs b/runtime/src/consensus/state/beacon.rs
index d91f19832f5..e71250df619 100644
--- a/runtime/src/consensus/state/beacon.rs
+++ b/runtime/src/consensus/state/beacon.rs
@@ -9,7 +9,7 @@ use crate::{
         state::StateError,
     },
     key_format,
-    storage::mkvs::ImmutableMKVS,
+    storage::mkvs::{FallibleMKVS, ImmutableMKVS},
 };
 
 /// Consensus beacon state wrapper.
@@ -65,6 +65,39 @@ impl<'a, T: ImmutableMKVS> ImmutableState<'a, T> {
     }
 }
 
+/// Mutable consensus beacon state wrapper.
+pub struct MutableState;
+
+impl MutableState {
+    /// Set current epoch state.
+    pub fn set_epoch_state<S: FallibleMKVS>(
+        mkvs: &mut S,
+        ctx: Context,
+        epoch_state: EpochTimeState,
+    ) -> Result<(), StateError> {
+        mkvs.insert(
+            ctx,
+            &CurrentEpochKeyFmt(()).encode(),
+            &cbor::to_vec(epoch_state),
+        )?;
+        Ok(())
+    }
+
+    /// Set future epoch state.
+    pub fn set_future_epoch_state<S: FallibleMKVS>(
+        mkvs: &mut S,
+        ctx: Context,
+        epoch_state: EpochTimeState,
+    ) -> Result<(), StateError> {
+        mkvs.insert(
+            ctx,
+            &FutureEpochKeyFmt(()).encode(),
+            &cbor::to_vec(epoch_state),
+        )?;
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
@@ -73,12 +106,58 @@ mod test {
         common::crypto::hash::Hash,
         storage::mkvs::{
             interop::{Fixture, ProtocolServer},
+            sync::NoopReadSyncer,
             Root, RootType, Tree,
         },
     };
 
     use super::*;
 
+    #[test]
+    fn test_mutable_state() {
+        let mut mkvs = Tree::builder()
+            .with_root_type(RootType::State)
+            .build(Box::new(NoopReadSyncer));
+
+        let ctx = Arc::new(Context::background());
+
+        MutableState::set_epoch_state(
+            &mut mkvs,
+            Context::create_child(&ctx),
+            EpochTimeState {
+                epoch: 10,
+                height: 100,
+            },
+        )
+        .unwrap();
+
+        MutableState::set_future_epoch_state(
+            &mut mkvs,
+            Context::create_child(&ctx),
+            EpochTimeState {
+                epoch: 11,
+                height: 110,
+            },
+        )
+        .unwrap();
+
+        let beacon_state = ImmutableState::new(&mkvs);
+
+        // Test current epoch state.
+        let epoch_state = beacon_state
+            .epoch_state(Context::create_child(&ctx))
+            .expect("epoch state query should work");
+        assert_eq!(10u64, epoch_state.epoch, "expected epoch should match");
+        assert_eq!(100i64, epoch_state.height, "expected height should match");
+
+        // Test future epoch state.
+        let epoch_state = beacon_state
+            .future_epoch_state(Context::create_child(&ctx))
+            .expect("future epoch state query should work");
+        assert_eq!(11u64, epoch_state.epoch, "expected epoch should match");
+        assert_eq!(110i64, epoch_state.height, "expected height should match");
+    }
+
     #[test]
     fn test_beacon_state_interop() {
         // Keep in sync with go/consensus/tendermint/apps/beacon/state/interop/interop.go.