CRC: move TestFlag related functions to seperate test modules

Signed-off-by: Jacinta Ferrant <[email protected]>

Author: jferrant

stacks-common/src/util/mod.rs

@@ -36,28 +36,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use std::{error, fmt, thread, time};
 
 #[cfg(any(test, feature = "testing"))]
-#[derive(Clone)]
-pub struct TestFlag<T>(pub std::sync::Arc<std::sync::Mutex<Option<T>>>);
-
-#[cfg(any(test, feature = "testing"))]
-impl<T: Default + Clone> Default for TestFlag<T> {
-    fn default() -> Self {
-        Self(std::sync::Arc::new(std::sync::Mutex::new(None)))
-    }
-}
-
-#[cfg(any(test, feature = "testing"))]
-impl<T: Default + Clone> TestFlag<T> {
-    /// Set the test flag to the given value
-    pub fn set(&self, value: T) {
-        *self.0.lock().unwrap() = Some(value);
-    }
-
-    /// Get the test flag value. Defaults otherwise.
-    pub fn get(&self) -> T {
-        self.0.lock().unwrap().clone().unwrap_or_default().clone()
-    }
-}
+pub mod tests;
 
 pub fn get_epoch_time_secs() -> u64 {
     let start = SystemTime::now();

stacks-common/src/util/tests.rs

@@ -0,0 +1,99 @@
+// Copyright (C) 2020-2024 Stacks Open Internet Foundation
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::sync::{Arc, Mutex};
+/// `TestFlag` is a thread-safe utility designed for managing shared state in testing scenarios. It wraps
+/// a value of type `T` inside an `Arc<Mutex<Option<T>>>`, allowing you to set and retrieve a value
+/// across different parts of your codebase while ensuring thread safety.
+///
+/// This structure is particularly useful when:
+/// - You need a global or static variable in tests.
+/// - You want to control the execution of custom test code paths by setting and checking a shared value.
+///
+/// # Type Parameter
+/// - `T`: The type of the value managed by the `TestFlag`. It must implement the `Default` and `Clone` traits.
+///
+/// # Examples
+///
+/// ```rust
+/// use stacks_common::util::tests::TestFlag;
+/// use std::sync::{Arc, Mutex};
+///
+/// // Create a TestFlag instance
+/// let test_flag = TestFlag::default();
+///
+/// // Set a value in the test flag
+/// test_flag.set("test_value".to_string());
+///
+/// // Retrieve the value
+/// assert_eq!(test_flag.get(), "test_value".to_string());
+///
+/// // Reset the value to default
+/// test_flag.set("".to_string());
+/// assert_eq!(test_flag.get(), "".to_string());
+/// ```
+#[derive(Clone)]
+pub struct TestFlag<T>(pub Arc<Mutex<Option<T>>>);
+
+impl<T: Default + Clone> Default for TestFlag<T> {
+    fn default() -> Self {
+        Self(Arc::new(Mutex::new(None)))
+    }
+}
+
+impl<T: Default + Clone> TestFlag<T> {
+    /// Sets the value of the test flag.
+    ///
+    /// This method updates the value stored inside the `TestFlag`, replacing any existing value.
+    ///
+    /// # Arguments
+    /// - `value`: The new value to set for the `TestFlag`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// let test_flag = TestFlag::default();
+    /// test_flag.set(42);
+    /// assert_eq!(test_flag.get(), 42);
+    /// ```
+    pub fn set(&self, value: T) {
+        *self.0.lock().unwrap() = Some(value);
+    }
+
+    /// Retrieves the current value of the test flag.
+    ///
+    /// If no value has been set, this method returns the default value for the type `T`.
+    ///
+    /// # Returns
+    /// - The current value of the test flag, or the default value of `T` if none has been set.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// let test_flag = TestFlag::default();
+    ///
+    /// // Get the default value
+    /// assert_eq!(test_flag.get(), 0); // For T = i32, default is 0
+    ///
+    /// // Set a value
+    /// test_flag.set(123);
+    ///
+    /// // Get the updated value
+    /// assert_eq!(test_flag.get(), 123);
+    /// ```
+    pub fn get(&self) -> T {
+        self.0.lock().unwrap().clone().unwrap_or_default().clone()
+    }
+}

stacks-signer/src/v0/mod.rs

@@ -17,6 +17,10 @@
 /// The signer module for processing events
 pub mod signer;
 
+#[cfg(any(test, feature = "testing"))]
+/// Test specific functions for the signer module
+pub mod tests;
+
 use libsigner::v0::messages::SignerMessage;
 
 use crate::v0::signer::Signer;

stacks-signer/src/v0/signer.rs

@@ -15,8 +15,6 @@
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::sync::mpsc::Sender;
-#[cfg(any(test, feature = "testing"))]
-use std::sync::LazyLock;
 use std::time::{Duration, Instant};
 
 use blockstack_lib::chainstate::nakamoto::{NakamotoBlock, NakamotoBlockHeader};
@@ -35,12 +33,8 @@ use libsigner::v0::messages::{
 use libsigner::{BlockProposal, SignerEvent};
 use slog::{slog_debug, slog_error, slog_info, slog_warn};
 use stacks_common::types::chainstate::StacksAddress;
-#[cfg(any(test, feature = "testing"))]
-use stacks_common::types::chainstate::StacksPublicKey;
 use stacks_common::util::get_epoch_time_secs;
 use stacks_common::util::secp256k1::MessageSignature;
-#[cfg(any(test, feature = "testing"))]
-use stacks_common::util::TestFlag;
 use stacks_common::{debug, error, info, warn};
 
 use crate::chainstate::{ProposalEvalConfig, SortitionsView};
@@ -50,27 +44,13 @@ use crate::runloop::SignerResult;
 use crate::signerdb::{BlockInfo, BlockState, SignerDb};
 use crate::Signer as SignerTrait;
 
-#[cfg(any(test, feature = "testing"))]
-/// A global variable that can be used to reject all block proposals if the signer's public key is in the provided list
-pub static TEST_REJECT_ALL_BLOCK_PROPOSAL: LazyLock<TestFlag<Vec<StacksPublicKey>>> =
-    LazyLock::new(TestFlag::default);
-
-#[cfg(any(test, feature = "testing"))]
-/// A global variable that can be used to ignore block proposals if the signer's public key is in the provided list
-pub static TEST_IGNORE_ALL_BLOCK_PROPOSALS: LazyLock<TestFlag<Vec<StacksPublicKey>>> =
-    LazyLock::new(TestFlag::default);
-
-#[cfg(any(test, feature = "testing"))]
-/// Pause the block broadcast
-pub static TEST_PAUSE_BLOCK_BROADCAST: LazyLock<TestFlag<bool>> = LazyLock::new(TestFlag::default);
-
-#[cfg(any(test, feature = "testing"))]
-/// Skip broadcasting the block to the network
-pub static TEST_SKIP_BLOCK_BROADCAST: LazyLock<TestFlag<bool>> = LazyLock::new(TestFlag::default);
-
 /// The stacks signer registered for the reward cycle
 #[derive(Debug)]
 pub struct Signer {
+    /// The private key of the signer
+    #[cfg(any(test, feature = "testing"))]
+    pub private_key: StacksPrivateKey,
+    #[cfg(not(any(test, feature = "testing")))]
     /// The private key of the signer
     private_key: StacksPrivateKey,
     /// The stackerdb client
@@ -175,20 +155,8 @@ impl SignerTrait<SignerMessage> for Signer {
                     match message {
                         SignerMessage::BlockProposal(block_proposal) => {
                             #[cfg(any(test, feature = "testing"))]
-                            {
-                                let public_keys = TEST_IGNORE_ALL_BLOCK_PROPOSALS.get();
-                                if public_keys.contains(
-                                    &stacks_common::types::chainstate::StacksPublicKey::from_private(
-                                        &self.private_key,
-                                    ),
-                                ) {
-                                        warn!("{self}: Ignoring block proposal due to testing directive";
-                                            "block_id" => %block_proposal.block.block_id(),
-                                            "height" => block_proposal.block.header.chain_length,
-                                            "consensus_hash" => %block_proposal.block.header.consensus_hash
-                                        );
-                                        continue;
-                                }
+                            if self.test_ignore_all_block_proposals(block_proposal) {
+                                continue;
                             }
                             self.handle_block_proposal(
                                 stacks_client,
@@ -1121,87 +1089,6 @@ impl Signer {
         }
     }
 
-    #[cfg(any(test, feature = "testing"))]
-    fn test_skip_block_broadcast(&self, block: &NakamotoBlock) -> bool {
-        if TEST_SKIP_BLOCK_BROADCAST.get() {
-            let block_hash = block.header.signer_signature_hash();
-            warn!(
-                "{self}: Skipping block broadcast due to testing directive";
-                "block_id" => %block.block_id(),
-                "height" => block.header.chain_length,
-                "consensus_hash" => %block.header.consensus_hash
-            );
-
-            if let Err(e) = self
-                .signer_db
-                .set_block_broadcasted(&block_hash, get_epoch_time_secs())
-            {
-                warn!("{self}: Failed to set block broadcasted for {block_hash}: {e:?}");
-            }
-            return true;
-        }
-        false
-    }
-
-    #[cfg(any(test, feature = "testing"))]
-    fn test_reject_block_proposal(
-        &mut self,
-        block_proposal: &BlockProposal,
-        block_info: &mut BlockInfo,
-        block_response: Option<BlockResponse>,
-    ) -> Option<BlockResponse> {
-        let public_keys = TEST_REJECT_ALL_BLOCK_PROPOSAL.get();
-        if public_keys.contains(
-            &stacks_common::types::chainstate::StacksPublicKey::from_private(&self.private_key),
-        ) {
-            warn!("{self}: Rejecting block proposal automatically due to testing directive";
-                "block_id" => %block_proposal.block.block_id(),
-                "height" => block_proposal.block.header.chain_length,
-                "consensus_hash" => %block_proposal.block.header.consensus_hash
-            );
-            if let Err(e) = block_info.mark_locally_rejected() {
-                warn!("{self}: Failed to mark block as locally rejected: {e:?}",);
-            };
-            // We must insert the block into the DB to prevent subsequent repeat proposals being accepted (should reject
-            // as invalid since we rejected in a prior round if this crops up again)
-            // in case this is the first time we saw this block. Safe to do since this is testing case only.
-            self.signer_db
-                .insert_block(block_info)
-                .unwrap_or_else(|e| self.handle_insert_block_error(e));
-            Some(BlockResponse::rejected(
-                block_proposal.block.header.signer_signature_hash(),
-                RejectCode::TestingDirective,
-                &self.private_key,
-                self.mainnet,
-                self.signer_db.calculate_tenure_extend_timestamp(
-                    self.proposal_config.tenure_idle_timeout,
-                    &block_proposal.block,
-                    false,
-                ),
-            ))
-        } else {
-            block_response
-        }
-    }
-
-    #[cfg(any(test, feature = "testing"))]
-    fn test_pause_block_broadcast(&self, block_info: &BlockInfo) {
-        if TEST_PAUSE_BLOCK_BROADCAST.get() {
-            // Do an extra check just so we don't log EVERY time.
-            warn!("{self}: Block broadcast is stalled due to testing directive.";
-                "block_id" => %block_info.block.block_id(),
-                "height" => block_info.block.header.chain_length,
-            );
-            while TEST_PAUSE_BLOCK_BROADCAST.get() {
-                std::thread::sleep(std::time::Duration::from_millis(10));
-            }
-            info!("{self}: Block validation is no longer stalled due to testing directive.";
-                "block_id" => %block_info.block.block_id(),
-                "height" => block_info.block.header.chain_length,
-            );
-        }
-    }
-
     /// Send a mock signature to stackerdb to prove we are still alive
     fn mock_sign(&mut self, mock_proposal: MockProposal) {
         info!("{self}: Mock signing mock proposal: {mock_proposal:?}");
@@ -1216,7 +1103,7 @@ impl Signer {
     }
 
     /// Helper for logging insert_block error
-    fn handle_insert_block_error(&self, e: DBError) {
+    pub fn handle_insert_block_error(&self, e: DBError) {
         error!("{self}: Failed to insert block into signer-db: {e:?}");
         panic!("{self} Failed to write block to signerdb: {e}");
     }