From cb072fa97aade0bbc5e289e6b13e6523512d5f77 Mon Sep 17 00:00:00 2001
From: Max Summe <maximilian.summe@dfinity.org>
Date: Thu, 24 Oct 2024 16:40:50 -0700
Subject: [PATCH] Fix more tests

---
 rs/nns/governance/src/governance.rs     |  4 ++--
 rs/nns/governance/src/neuron_store.rs   | 16 +++++++++++++---
 rs/nns/governance/tests/fixtures/mod.rs |  2 +-
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/rs/nns/governance/src/governance.rs b/rs/nns/governance/src/governance.rs
index 37e7368fb0f..15981d79ae5 100644
--- a/rs/nns/governance/src/governance.rs
+++ b/rs/nns/governance/src/governance.rs
@@ -1907,8 +1907,8 @@ impl Governance {
         )
     }
 
-    pub fn clone_proto(&self) -> GovernanceProto {
-        let neurons = self.neuron_store.clone_neurons();
+    pub fn __get_state_for_test(&self) -> GovernanceProto {
+        let neurons = self.neuron_store.__get_neurons_for_tests();
         let heap_topic_followee_index = self.neuron_store.clone_topic_followee_index();
         let heap_governance_proto = self.heap_data.clone();
         let rng_seed = self.env.get_rng_seed();
diff --git a/rs/nns/governance/src/neuron_store.rs b/rs/nns/governance/src/neuron_store.rs
index 4a4bc8b7e7e..5e48adb307d 100644
--- a/rs/nns/governance/src/neuron_store.rs
+++ b/rs/nns/governance/src/neuron_store.rs
@@ -429,11 +429,21 @@ impl NeuronStore {
 
     /// Clones all the neurons. This is only used for testing.
     /// TODO(NNS-2474) clean it up after NNSState stop using GovernanceProto.
-    pub fn clone_neurons(&self) -> BTreeMap<u64, NeuronProto> {
-        self.heap_neurons
+    pub fn __get_neurons_for_tests(&self) -> BTreeMap<u64, NeuronProto> {
+        let mut stable_neurons = with_stable_neuron_store(|stable_store| {
+            stable_store
+                .range_neurons(..)
+                .map(|neuron| (neuron.id().id, neuron.into()))
+                .collect::<BTreeMap<u64, NeuronProto>>()
+        });
+        let mut heap_neurons = self
+            .heap_neurons
             .iter()
             .map(|(id, neuron)| (*id, neuron.clone().into()))
-            .collect()
+            .collect::<BTreeMap<u64, NeuronProto>>();
+
+        stable_neurons.extend(heap_neurons);
+        stable_neurons
     }
 
     pub fn clone_topic_followee_index(&self) -> HashMap<i32, FollowersMap> {
diff --git a/rs/nns/governance/tests/fixtures/mod.rs b/rs/nns/governance/tests/fixtures/mod.rs
index 304a12aade9..32578ef303f 100755
--- a/rs/nns/governance/tests/fixtures/mod.rs
+++ b/rs/nns/governance/tests/fixtures/mod.rs
@@ -734,7 +734,7 @@ impl NNS {
             .ledger
             .accounts
             .clone();
-        let governance_proto = self.governance.clone_proto();
+        let governance_proto = self.governance.__get_state_for_test();
         NNSState {
             now: self.now(),
             accounts,