tower-test

pub fn pair<T, U>() -> (Mock<T, U>, Handle<T, U>) {
    let (tx, rx) = mpsc::unbounded_channel();
    let tx = Mutex::new(tx);

    let state = Arc::new(Mutex::new(State::new()));

    let mock = Mock {
        id: 0,
        tx,
        state: state.clone(),
        can_send: false,
    };

    let handle = Handle { rx, state };

    (mock, handle)
}

struct State {
    /// Tracks the number of requests that can be sent through
    rem: u64,

    /// Tasks that are blocked
    tasks: HashMap<u64, Waker>,

    /// Tracks if the `Handle` dropped
    is_closed: bool,

    /// Tracks the ID for the next mock clone
    next_clone_id: u64,

    /// Tracks the next error to yield (if any)
    err_with: Option<Error>,
}

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let mut state = self.state.lock().unwrap();

        if state.is_closed {
            return Poll::Ready(Err(error::Closed::new().into()));
        }

        if let Some(e) = state.err_with.take() {
            return Poll::Ready(Err(e));
        }

        if self.can_send {
            return Poll::Ready(Ok(()));
        }

        if state.rem > 0 {
            assert!(!state.tasks.contains_key(&self.id));

            // Returning `Ready` means the next call to `call` must succeed.
            self.can_send = true;

            Poll::Ready(Ok(()))
        } else {
            // Bit weird... but whatevz
            *state
                .tasks
                .entry(self.id)
                .or_insert_with(|| cx.waker().clone()) = cx.waker().clone();

            Poll::Pending
        }
    }

    pub fn allow(&mut self, num: u64) {
        let mut state = self.state.lock().unwrap();
        state.rem = num;

        if num > 0 {
            for (_, task) in state.tasks.drain() {
                task.wake();
            }
        }
    }

    /// Make the next poll_ method error with the given error.
    pub fn send_error<E: Into<Error>>(&mut self, e: E) {
        let mut state = self.state.lock().unwrap();
        state.err_with = Some(e.into());

        for (_, task) in state.tasks.drain() {
            task.wake();
        }
    }

tower-test = "0.4"
tokio-test = "0.4"
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use std::panic::{catch_unwind, AssertUnwindSafe};
use tokio_test::task;
use tower_test::mock;

fn show<T: std::fmt::Debug, E: ToString>(p: std::task::Poll<Result<T, E>>) -> String {
  format!("{:?}", p.map(|r| r.map_err(|e| e.to_string())))
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let (mut svc, mut handle) = mock::spawn::<&'static str, &'static str>();
  let (mut other, other_handle) = mock::spawn::<&'static str, &'static str>();

  println!("1. no allow() yet");
  println!("   poll_ready -> {}", show(svc.poll_ready()));
  let mut fut = task::spawn(svc.call("a"));
  let (req, send) = handle.next_request().await.unwrap();
  send.send_response("A");
  println!("   got {req:?}, response {}", show(fut.poll()));

  println!("2. allow(0), then allow(1)");
  handle.allow(0);
  println!("   poll_ready -> {}", show(svc.poll_ready()));
  handle.allow(1);
  println!("   woken: {}, poll_ready -> {}", svc.is_woken(), show(svc.poll_ready()));
  println!("   poll_ready again before call -> {}", show(svc.poll_ready()));
  let _f = svc.call("b");
  println!("   after call, poll_ready -> {}", show(svc.poll_ready()));

  println!("3. allow sets, it does not add");
  handle.allow(2);
  handle.allow(1);
  println!("   poll_ready -> {}", show(svc.poll_ready()));
  let _f = svc.call("c");
  println!("   poll_ready -> {}", show(svc.poll_ready()));

  println!("4. allow(1) and two clones");
  let mut c1 = mock::Spawn::new(svc.get_ref().clone());
  let mut c2 = mock::Spawn::new(svc.get_ref().clone());
  handle.allow(1);
  println!("   c1 -> {}, c2 -> {}", show(c1.poll_ready()), show(c2.poll_ready()));
  let _f1 = c1.call("d1");
  let _f2 = c2.call("d2");
  println!("   both called; c1 -> {}", show(c1.poll_ready()));

  println!("5. Handle::send_error hits the next poll_ready once");
  handle.allow(1);
  let mut c3 = mock::Spawn::new(svc.get_ref().clone());
  handle.send_error("overloaded");
  println!("   c3 -> {}", show(c3.poll_ready()));
  println!("   c3 -> {}", show(c3.poll_ready()));

  println!("6. dropping a SendResponse");
  let mut fut = task::spawn(c3.call("e"));
  let send = loop {
    let (req, send) = handle.next_request().await.unwrap();
    if req == "e" {
      break send;
    }
  };
  drop(send);
  println!("   response {}", show(fut.poll()));

  println!("7. after the Handle is dropped");
  drop(handle);
  let mut c5 = mock::Spawn::new(svc.get_ref().clone());
  println!("   poll_ready -> {}", show(c5.poll_ready()));
  let mut fut = task::spawn(c5.call("g"));
  println!("   call without ready -> {}", show(fut.poll()));

  println!("8. call without poll_ready (a separate pair)");
  std::panic::set_hook(Box::new(|_| {})); // keep the default hook quiet
  let r = catch_unwind(AssertUnwindSafe(|| {
    let _ = other.call("f");
  }));
  let p = r.unwrap_err();
  println!("   panic: {:?}", p.downcast_ref::<&str>());

  let r = catch_unwind(AssertUnwindSafe(|| other.poll_ready().is_ready()));
  let p = r.unwrap_err();
  println!("   then poll_ready panics: {:?}", p.downcast_ref::<String>());

  let r = catch_unwind(AssertUnwindSafe(|| drop(other_handle)));
  println!("   dropping its Handle panics: {:?}", r.unwrap_err().downcast_ref::<String>());
  let r = catch_unwind(AssertUnwindSafe(|| drop(other)));
  println!("   dropping the Mock panics: {:?}", r.unwrap_err().downcast_ref::<String>());
}

    fn call(&mut self, request: T) -> Self::Future {
        // Make sure that the service has capacity
        let mut state = self.state.lock().unwrap();

        if state.is_closed {
            return ResponseFuture::closed();
        }

        if !self.can_send {
            panic!("service not ready; poll_ready must be called first");
        }

        self.can_send = false;

        // Decrement the number of remaining requests that can be sent
        if state.rem > 0 {
            state.rem -= 1;
        }
        // ...

    /// Asynchronously gets the next request
    pub fn poll_request(&mut self) -> Poll<Option<Request<T, U>>> {
        tokio_test::task::spawn(()).enter(|cx, _| Box::pin(self.rx.recv()).as_mut().poll(cx))
    }

    /// Gets the next request.
    pub async fn next_request(&mut self) -> Option<Request<T, U>> {
        self.rx.recv().await
    }

    ($mock_handle:expr, $expect:expr, $($arg:tt)*) => {{
        let (actual, send_response) = match $mock_handle.next_request().await {
            Some(r) => r,
            None => panic!("expected a request but none was received."),
        };
        // ...
        {
            assert_eq!(actual, $expect, $($arg)*);
        }
        send_response
    }};

pub struct Spawn<T> {
    inner: T,
    task: task::Spawn<()>,
}

tower = { version = "=0.5.3", features = ["limit"] }
tower-test = "0.4"
tokio-test = "0.4"
tokio = { version = "=1.54.1", features = ["macros", "rt"] }

use tokio_test::{assert_pending, assert_ready, assert_ready_ok, task};
use tower::limit::ConcurrencyLimitLayer;
use tower_test::{assert_request_eq, mock};

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let (mut svc, mut handle) = mock::spawn_layer::<&str, &str, _>(ConcurrencyLimitLayer::new(1));

  assert_ready_ok!(svc.poll_ready());
  let mut first = task::spawn(svc.call("one"));
  println!("first called");

  // the only permit is inside `first`, so the layer must push back
  assert_pending!(svc.poll_ready());
  println!("second poll_ready: Pending, woken = {}", svc.is_woken());

  assert_request_eq!(handle, "one").send_response("1");
  assert_eq!(assert_ready!(first.poll()).unwrap(), "1");
  println!("first done, woken = {}", svc.is_woken());
  assert_pending!(svc.poll_ready()); // permit is still in `first`

  drop(first);
  println!("first dropped, woken = {}", svc.is_woken());
  assert_ready_ok!(svc.poll_ready());
  println!("second poll_ready: Ready");
}

first called
second poll_ready: Pending, woken = false
first done, woken = false
first dropped, woken = true
second poll_ready: Ready

1. no allow() yet
   poll_ready -> Ready(Ok(()))
   got "a", response Ready(Ok("A"))
2. allow(0), then allow(1)
   poll_ready -> Pending
   woken: true, poll_ready -> Ready(Ok(()))
   poll_ready again before call -> Ready(Ok(()))
   after call, poll_ready -> Pending
3. allow sets, it does not add
   poll_ready -> Ready(Ok(()))
   poll_ready -> Pending
4. allow(1) and two clones
   c1 -> Ready(Ok(())), c2 -> Ready(Ok(()))
   both called; c1 -> Pending
5. Handle::send_error hits the next poll_ready once
   c3 -> Ready(Err("overloaded"))
   c3 -> Ready(Ok(()))
6. dropping a SendResponse
   response Ready(Err("service closed"))
7. after the Handle is dropped
   poll_ready -> Ready(Err("service closed"))
   call without ready -> Ready(Err("service closed"))
8. call without poll_ready (a separate pair)
   panic: Some("service not ready; poll_ready must be called first")
   then poll_ready panics: Some("called `Result::unwrap()` on an `Err` value: PoisonError { .. }")
   dropping its Handle panics: Some("PoisonError { .. }")
   dropping the Mock panics: Some("PoisonError { .. }")