blueprint_eigenlayer_testing_utils/
runner.rs

1#![allow(dead_code)]
2
3use blueprint_core::Job;
4use blueprint_core_testing_utils::runner::{TestEnv, TestRunner};
5use blueprint_runner::BackgroundService;
6use blueprint_runner::config::BlueprintEnvironment;
7use blueprint_runner::eigenlayer::bls::EigenlayerBLSConfig;
8use blueprint_runner::error::{JobCallError, RunnerError as Error};
9use tokio::sync::Mutex;
10use tokio::task::JoinHandle;
11
12pub struct EigenlayerBLSTestEnv<Ctx> {
13    pub runner: Option<TestRunner<Ctx>>,
14    pub config: EigenlayerBLSConfig,
15    pub env: BlueprintEnvironment,
16    pub runner_handle: Mutex<Option<JoinHandle<Result<(), Error>>>>,
17}
18
19impl<Ctx> TestEnv for EigenlayerBLSTestEnv<Ctx>
20where
21    Ctx: Clone + Send + Sync + 'static,
22{
23    type Config = EigenlayerBLSConfig;
24    type Context = Ctx;
25
26    fn new(config: Self::Config, env: BlueprintEnvironment) -> Result<Self, Error> {
27        let runner = TestRunner::new(config, env.clone());
28
29        Ok(Self {
30            runner: Some(runner),
31            config,
32            env,
33            runner_handle: Mutex::new(None),
34        })
35    }
36
37    fn add_job<J, T>(&mut self, job: J)
38    where
39        J: Job<T, Self::Context> + Send + Sync + 'static,
40        T: 'static,
41    {
42        self.runner
43            .as_mut()
44            .expect("Runner already running")
45            .add_job(job);
46    }
47
48    fn add_background_service<B>(&mut self, service: B)
49    where
50        B: BackgroundService + Send + 'static,
51    {
52        self.runner
53            .as_mut()
54            .expect("Runner already running")
55            .add_background_service(service);
56    }
57
58    fn get_blueprint_config(&self) -> BlueprintEnvironment {
59        self.env.clone()
60    }
61
62    async fn run_runner(&mut self, context: Self::Context) -> Result<(), Error> {
63        // Spawn the runner in a background task
64        let runner = self.runner.take().expect("Runner already running");
65        let handle = tokio::spawn(async move { runner.run(context).await });
66
67        let mut guard = self.runner_handle.lock().await;
68        *guard = Some(handle);
69
70        // Brief delay to allow for startup
71        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
72
73        // Just check if it failed immediately
74        let handle = guard.take().expect("value was just set");
75        if !handle.is_finished() {
76            // Put the handle back since the runner is still running
77            *guard = Some(handle);
78            blueprint_core::info!("Runner started successfully");
79            return Ok(());
80        }
81
82        blueprint_core::info!("Runner task finished");
83        match handle.await {
84            Ok(Ok(())) => Ok(()),
85            Ok(Err(e)) => {
86                blueprint_core::error!("Runner failed during startup: {}", e);
87                Err(e)
88            }
89            Err(e) => {
90                blueprint_core::error!("Runner task panicked: {}", e);
91                Err(JobCallError::JobDidntFinish(e).into())
92            }
93        }
94    }
95}
96
97impl<Ctx> Drop for EigenlayerBLSTestEnv<Ctx> {
98    fn drop(&mut self) {
99        futures::executor::block_on(async {
100            let mut guard = self.runner_handle.lock().await;
101            if let Some(handle) = guard.take() {
102                handle.abort();
103            }
104        });
105    }
106}