Skip to content

Commit

Permalink
feat: add Test traits for Header and Block
Browse files Browse the repository at this point in the history
  • Loading branch information
mattsse committed Jan 8, 2025
1 parent 28d5231 commit f7ea2f4
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 22 deletions.
26 changes: 26 additions & 0 deletions crates/primitives-traits/src/block/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,3 +88,29 @@ where
(self.header, self.body)
}
}

/// An extension trait for [`Block`]s that allows for mutable access to the block's internals.
///
/// This allows for modifying the block's header and body for testing purposes.
#[cfg(any(test, feature = "test-utils"))]
pub trait TestBlock: Block {
/// Returns mutable reference to block body.
fn body_mut(&mut self) -> &mut Self::Body;

/// Returns mutable reference to block header.
fn header_mut(&mut self) -> &mut Self::Header;
}

#[cfg(any(test, feature = "test-utils"))]
impl<T> TestBlock for alloy_consensus::Block<T>
where
T: SignedTransaction,
{
fn body_mut(&mut self) -> &mut Self::Body {
&mut self.body
}

fn header_mut(&mut self) -> &mut Self::Header {
&mut self.header
}
}
36 changes: 18 additions & 18 deletions crates/primitives-traits/src/header/sealed.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,24 @@ impl Decodable for SealedHeader {
}
}

impl<H> From<SealedHeader<H>> for Sealed<H> {
fn from(value: SealedHeader<H>) -> Self {
Self::new_unchecked(value.header, value.hash)
}
}

#[cfg(any(test, feature = "arbitrary"))]
impl<'a, H> arbitrary::Arbitrary<'a> for SealedHeader<H>
where
H: for<'b> arbitrary::Arbitrary<'b> + Sealable,
{
fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
let header = H::arbitrary(u)?;

Ok(Self::seal(header))
}
}

#[cfg(any(test, feature = "test-utils"))]
impl SealedHeader {
/// Updates the block header.
Expand Down Expand Up @@ -152,24 +170,6 @@ impl SealedHeader {
}
}

impl<H> From<SealedHeader<H>> for Sealed<H> {
fn from(value: SealedHeader<H>) -> Self {
Self::new_unchecked(value.header, value.hash)
}
}

#[cfg(any(test, feature = "arbitrary"))]
impl<'a, H> arbitrary::Arbitrary<'a> for SealedHeader<H>
where
H: for<'b> arbitrary::Arbitrary<'b> + Sealable,
{
fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
let header = H::arbitrary(u)?;

Ok(Self::seal(header))
}
}

/// Bincode-compatible [`SealedHeader`] serde implementation.
#[cfg(feature = "serde-bincode-compat")]
pub(super) mod serde_bincode_compat {
Expand Down
39 changes: 37 additions & 2 deletions crates/primitives-traits/src/header/test_utils.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,45 @@
//! Test utilities to generate random valid headers.
//! Test utilities for the block header.
use alloy_consensus::Header;
use alloy_primitives::B256;
use alloy_primitives::{BlockHash, BlockNumber, B256, U256};
use proptest::{arbitrary::any, prop_compose};
use proptest_arbitrary_interop::arb;

/// A helper trait for [`Header`]s that allows for mutable access to the headers values.
///
/// This allows for modifying the header for testing purposes.
pub trait TestHeader {
/// Updates the parent block hash.
fn set_parent_hash(&mut self, hash: BlockHash);

/// Updates the block number.
fn set_block_number(&mut self, number: BlockNumber);

/// Updates the block state root.
fn set_state_root(&mut self, state_root: B256);

/// Updates the block difficulty.
fn set_difficulty(&mut self, difficulty: U256);
}

impl TestHeader for Header {
fn set_parent_hash(&mut self, hash: BlockHash) {
self.parent_hash = hash
}

fn set_block_number(&mut self, number: BlockNumber) {
self.number = number;
}

fn set_state_root(&mut self, state_root: B256) {
self.state_root = state_root;
}

fn set_difficulty(&mut self, difficulty: U256) {
self.difficulty = difficulty;
}
}

/// Generates a header which is valid __with respect to past and future forks__. This means, for
/// example, that if the withdrawals root is present, the base fee per gas is also present.
///
Expand Down
8 changes: 6 additions & 2 deletions crates/primitives-traits/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,6 @@ pub mod sync;

/// Common header types
pub mod header;
#[cfg(any(test, feature = "arbitrary", feature = "test-utils"))]
pub use header::test_utils;
pub use header::{Header, HeaderError, SealedHeader};

/// Bincode-compatible serde implementations for common abstracted types in Reth.
Expand Down Expand Up @@ -135,3 +133,9 @@ pub trait MaybeSerdeBincodeCompat {}
impl<T> MaybeSerdeBincodeCompat for T where T: crate::serde_bincode_compat::SerdeBincodeCompat {}
#[cfg(not(feature = "serde-bincode-compat"))]
impl<T> MaybeSerdeBincodeCompat for T {}

/// Utilities for testing.
#[cfg(any(test, feature = "test-utils"))]
pub mod test_utils {
pub use crate::{block::TestBlock, header::test_utils::*};
}

0 comments on commit f7ea2f4

Please sign in to comment.