Use a stash_ops helper

massalabs · Feb 9, 2023 · d928f4b · d928f4b
1 parent efd7614
commit d928f4b
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 13 deletions.
diff --git a/massa-pool-worker/src/tests/operation_pool_tests.rs b/massa-pool-worker/src/tests/operation_pool_tests.rs
@@ -19,9 +19,9 @@
 //!
 use crate::tests::tools::OpGenerator;
 
-use super::tools::{create_some_operations, operation_pool_test, pool_test};
+use super::tools::{create_some_operations, operation_pool_test, pool_test, stash_ops};
 use massa_execution_exports::test_exports::MockExecutionControllerMessage;
-use massa_models::{amount::Amount, prehash::PreHashMap, slot::Slot};
+use massa_models::{amount::Amount, slot::Slot};
 use massa_pool_exports::PoolConfig;
 use std::time::Duration;
 
@@ -61,18 +61,19 @@ fn test_pool() {
             let mut thread_tx_lists = vec![Vec::new(); pool_config.thread_count as usize];
             for i in 0..18 {
                 // setup operation
-                let fee = 40 + i;
                 let expire_period: u64 = 40 + i;
                 let start_period =
                     expire_period.saturating_sub(pool_config.operation_validity_periods);
                 let op = OpGenerator::default()
                     .expirery(expire_period)
-                    .fee(Amount::from_raw(fee))
-                    .generate(); //get_transaction(expire_period, fee);
+                    .fee(Amount::from_raw(40 + i))
+                    .generate();
 
-                let mut storage = storage_base.clone_without_refs();
-                storage.store_operations(vec![op.clone()]);
-                pool.add_operations(storage);
+                stash_ops(
+                    vec![op.clone()],
+                    storage_base.clone_without_refs(),
+                    &mut pool,
+                );
 
                 //TODO: compare
                 // assert_eq!(storage.get_op_refs(), &Set::<OperationId>::default());

diff --git a/massa-pool-worker/src/tests/scenario.rs b/massa-pool-worker/src/tests/scenario.rs
@@ -17,15 +17,20 @@ use std::time::Duration;
 
 use crate::tests::tools::create_some_operations;
 use crate::tests::tools::pool_test;
+use crate::tests::tools::stash_ops;
 use crate::tests::tools::OpGenerator;
 use massa_execution_exports::test_exports::MockExecutionControllerMessage as ControllerMsg;
 use massa_models::address::Address;
 use massa_models::amount::Amount;
+use massa_models::operation::Operation;
 use massa_models::operation::OperationId;
 use massa_models::prehash::PreHashSet;
+use massa_models::secure_share::SecureShare;
 use massa_models::slot::Slot;
 use massa_pool_exports::PoolConfig;
+use massa_pool_exports::PoolController;
 use massa_signature::KeyPair;
+use massa_storage::Storage;
 
 /// # Test simple get operation
 /// Just try to get some operations stored in pool
@@ -57,9 +62,11 @@ fn test_simple_get_operations() {
             let creator_address = Address::from_public_key(&keypair.get_public_key());
             let creator_thread = creator_address.get_thread(config.thread_count);
 
-            // Setup the operations...
-            storage.store_operations(create_some_operations(10, &op_gen));
-            pool_controller.add_operations(storage.clone());
+            stash_ops(
+                create_some_operations(10, &op_gen),
+                storage.clone(),
+                &mut pool_controller,
+            );
 
             // Start mock execution thread.
             // Provides the data for `pool_controller.get_block_operations`

diff --git a/massa-pool-worker/src/tests/tools.rs b/massa-pool-worker/src/tests/tools.rs
@@ -10,8 +10,8 @@ use massa_models::{
     amount::Amount,
     block_id::BlockId,
     endorsement::{Endorsement, EndorsementSerializer, SecureShareEndorsement},
-    operation::{Operation, OperationSerializer, OperationType, SecureShareOperation},
-    secure_share::SecureShareContent,
+    operation::{Operation, OperationId, OperationSerializer, OperationType, SecureShareOperation},
+    secure_share::{SecureShare, SecureShareContent},
     slot::Slot,
 };
 use massa_pool_exports::{PoolChannels, PoolConfig, PoolController, PoolManager};
@@ -82,6 +82,16 @@ pub(crate) fn create_some_operations(n: usize, op_gen: &OpGenerator) -> Vec<Secu
     (0..n).map(|_| op_gen.generate()).collect()
 }
 
+/// Sets up storage and pool to contain the operation
+pub(crate) fn stash_ops(
+    ops: Vec<SecureShare<Operation, OperationId>>,
+    mut storage: Storage,
+    pool: &mut Box<dyn PoolController>,
+) {
+    storage.store_operations(ops);
+    pool.add_operations(storage);
+}
+
 /// Creates module mocks, providing the environment needed to run the provided closure
 pub fn pool_test<F>(cfg: PoolConfig, test: F)
 where