test(exex): improve PollOnce trait error handling (#8653)

Cargo.lock

@@ -6940,6 +6940,7 @@ dependencies = [
  "reth-provider",
  "reth-tasks",
  "reth-transaction-pool",
+ "thiserror",
  "tokio",
 ]
 

crates/exex/test-utils/Cargo.toml

@@ -36,3 +36,4 @@ tokio.workspace = true
 ## misc
 eyre.workspace = true
 rand.workspace = true
+thiserror.workspace = true

crates/exex/test-utils/src/lib.rs

@@ -41,6 +41,7 @@ use std::{
     sync::Arc,
     task::Poll,
 };
+use thiserror::Error;
 use tokio::sync::mpsc::{Sender, UnboundedReceiver};
 
 /// A test [`PoolBuilder`] that builds a [`TestPool`].
@@ -240,21 +241,36 @@ pub async fn test_exex_context() -> eyre::Result<(ExExContext<Adapter>, TestExEx
 
 /// An extension trait for polling an Execution Extension future.
 pub trait PollOnce {
-    /// Polls the given Execution Extension future __once__ and asserts that it is
-    /// [`Poll::Pending`]. The future should be (pinned)[`std::pin::pin`].
+    /// Polls the given Execution Extension future __once__. The future should be
+    /// (pinned)[`std::pin::pin`].
     ///
-    /// # Panics
-    /// If the future returns [`Poll::Ready`], because Execution Extension future should never
-    /// resolve.
-    fn poll_once(&mut self) -> impl Future<Output = ()> + Send;
+    /// # Returns
+    /// - `Ok(())` if the future returned [`Poll::Pending`]. The future can be polled again.
+    /// - `Err(PollOnceError::FutureIsReady)` if the future returned [`Poll::Ready`] without an
+    ///   error. The future should never resolve.
+    /// - `Err(PollOnceError::FutureError(err))` if the future returned [`Poll::Ready`] with an
+    ///   error. Something went wrong.
+    fn poll_once(&mut self) -> impl Future<Output = Result<(), PollOnceError>> + Send;
+}
+
+/// An Execution Extension future polling error.
+#[derive(Error, Debug)]
+pub enum PollOnceError {
+    /// The future returned [`Poll::Ready`] without an error, but it should never resolve.
+    #[error("Execution Extension future returned Ready, but it should never resolve")]
+    FutureIsReady,
+    /// The future returned [`Poll::Ready`] with an error.
+    #[error(transparent)]
+    FutureError(#[from] eyre::Error),
 }
 
 impl<F: Future<Output = eyre::Result<()>> + Unpin + Send> PollOnce for F {
-    async fn poll_once(&mut self) {
-        poll_fn(|cx| {
-            assert!(self.poll_unpin(cx).is_pending());
-            Poll::Ready(())
+    async fn poll_once(&mut self) -> Result<(), PollOnceError> {
+        poll_fn(|cx| match self.poll_unpin(cx) {
+            Poll::Ready(Ok(())) => Poll::Ready(Err(PollOnceError::FutureIsReady)),
+            Poll::Ready(Err(err)) => Poll::Ready(Err(PollOnceError::FutureError(err))),
+            Poll::Pending => Poll::Ready(Ok(())),
         })
-        .await;
+        .await
     }
 }

examples/exex/minimal/src/main.rs

@@ -60,7 +60,7 @@ mod tests {
     use reth_exex_test_utils::{test_exex_context, PollOnce};
 
     #[tokio::test]
-    async fn exex() -> eyre::Result<()> {
+    async fn test_exex() -> eyre::Result<()> {
         // Initialize a test Execution Extension context with all dependencies
         let (ctx, mut handle) = test_exex_context().await?;
 
@@ -83,7 +83,7 @@ mod tests {
         handle.assert_events_empty();
 
         // Poll the Execution Extension once to process incoming notifications
-        exex.poll_once().await;
+        exex.poll_once().await?;
 
         // Check that the Execution Extension emitted a `FinishedHeight` event with the correct
         // height