hotshot_testing/
test_launcher.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7use std::{collections::HashMap, marker::PhantomData, rc::Rc, sync::Arc};
8
9use hotshot::traits::{NodeImplementation, TestableNodeImplementation};
10use hotshot_example_types::storage_types::TestStorage;
11use hotshot_types::{
12    traits::{
13        network::{AsyncGenerator, ConnectedNetwork},
14        node_implementation::{NodeType, Versions},
15    },
16    HotShotConfig, ValidatorConfig,
17};
18
19use super::{test_builder::TestDescription, test_runner::TestRunner};
20use crate::test_task::TestTaskStateSeed;
21
22/// A type alias to help readability
23pub type Network<TYPES, I> = Arc<<I as NodeImplementation<TYPES>>::Network>;
24
25/// Wrapper for a function that takes a `node_id` and returns an instance of `T`.
26pub type Generator<T> = Rc<dyn Fn(u64) -> T>;
27
28/// generators for resources used by each node
29pub struct ResourceGenerators<TYPES: NodeType, I: TestableNodeImplementation<TYPES>> {
30    /// generate channels
31    pub channel_generator: AsyncGenerator<Network<TYPES, I>>,
32    /// generate new storage for each node
33    pub storage: Generator<TestStorage<TYPES>>,
34    /// configuration used to generate each hotshot node
35    pub hotshot_config: Generator<HotShotConfig<TYPES>>,
36    /// config that contains the signature keys
37    pub validator_config: Generator<ValidatorConfig<TYPES>>,
38}
39
40/// test launcher
41pub struct TestLauncher<TYPES: NodeType, I: TestableNodeImplementation<TYPES>, V: Versions> {
42    /// generator for resources
43    pub resource_generators: ResourceGenerators<TYPES, I>,
44    /// metadata used for tasks
45    pub metadata: TestDescription<TYPES, I, V>,
46    /// any additional test tasks to run
47    pub additional_test_tasks: Vec<Box<dyn TestTaskStateSeed<TYPES, I, V>>>,
48}
49
50impl<TYPES: NodeType, I: TestableNodeImplementation<TYPES>, V: Versions> TestLauncher<TYPES, I, V> {
51    /// launch the test
52    #[must_use]
53    pub fn launch<N: ConnectedNetwork<TYPES::SignatureKey>>(self) -> TestRunner<TYPES, I, V, N> {
54        TestRunner::<TYPES, I, V, N> {
55            launcher: self,
56            nodes: Vec::new(),
57            late_start: HashMap::new(),
58            next_node_id: 0,
59            _pd: PhantomData,
60        }
61    }
62    /// Modifies the config used when generating nodes with `f`
63    #[must_use]
64    pub fn map_hotshot_config(mut self, f: impl Fn(&mut HotShotConfig<TYPES>) + 'static) -> Self {
65        let mut test_config = self.metadata.test_config.clone();
66        f(&mut test_config);
67
68        let hotshot_config_generator = self.resource_generators.hotshot_config.clone();
69        let hotshot_config: Generator<_> = Rc::new(move |node_id| {
70            let mut result = (hotshot_config_generator)(node_id);
71            f(&mut result);
72
73            result
74        });
75
76        self.metadata.test_config = test_config;
77        self.resource_generators.hotshot_config = hotshot_config;
78
79        self
80    }
81}
hotshot_testing/test_launcher.rs

hotshot_testing/
test_launcher.rs
