Difficult to unit test workflows

[BrynCooke]

I'm using restate for the first time and the project looks great, but the ergonomics for unit testing could be improved.

There are two main issues:

1. When defining a trait for a workflow or a service self must not be present.

#[restate_sdk::workflow]
trait MyWorkflow {
  fn do_something(); // No self
}

This means that regular mocking libraries cannot be used with these traits. The restate macro adds self which honestly is confusing. Why not just require the trait functions to take self rather than add it magically?

2. Contex cannot be constructed and therefore any function that takes context cannot be unit tested. I want to test failure of my workflows for error handling, testcontainers is not a good fit as I will need to be able to set expectations and failures in a fine grained way.

It'd be great if Restate could provide unit test support for Context. The only solution I can see right now is to abstract Restate apis away completely in my codebase.

[slinkydeveloper]

Hey there, will answer point by point:

This means that regular mocking libraries cannot be used with these traits. The restate macro adds self which honestly is confusing. Why not just require the trait functions to take self rather than add it magically?

There is an ongoing discussion about the macro shape here #67, indeed some things might be too magical in the current macro and we might want to change it in the future, to simplify mocking for example.

Contex cannot be constructed and therefore any function that takes context cannot be unit tested.

The Context interface implements few traits, see https://docs.rs/restate-sdk/latest/restate_sdk/context/index.html#traits, one thing you can try to do is to reorganize your code in functions that take these traits individually (e.g. impl restate_sdk::prelude::ContextClient) and mock them. Perhaps one additional thing we can support in the macro is to define impl ContextClient etc as context argument (though honestly i don't know if this is possible...).

Providing a mock of the Context on our end is somewhat complicated, there was a discussion here restatedev/documentation#377 with more details.

[akkie]

There is an ongoing discussion about the macro shape here #67, indeed some things might be too magical in the current macro and we might want to change it in the future, to simplify mocking for example.

Doesn't the proposed implementation make things even more magical?

Providing a mock of the Context on our end is somewhat complicated, there was a discussion here restatedev/documentation#377 with more details.

Currently you can do something like this:

#[automock]
pub trait ObjectContextProvider {
    fn key(&self) -> String;
}

impl ObjectContextProvider for ObjectContext<'_> {
    fn key(&self) -> String {
        self.key().to_owned()
    }
}

#[restate_sdk::object]
pub trait MediaTextObject {
    async fn save(request: Json<SaveTextRequest>) -> Result<Json<SaveTextResponse>, HandlerError>;
}

pub struct MediaTextObjectImpl(pub Arc<dyn AppState>);

impl MediaTextObject for MediaTextObjectImpl {
    async fn save(
        &self,
        ctx: ObjectContext<'_>,
        request: Json<SaveTextRequest>,
    ) -> Result<Json<SaveTextResponse>, HandlerError> {
        self.save_internal(&ctx, request).await
    }
}

impl MediaTextObjectImpl {
    async fn save_internal<C: ObjectContextProvider>(
        &self,
        ctx: &C,
        Json(request): Json<SaveTextRequest>,
    ) -> Result<Json<SaveTextResponse>, HandlerError> {
        let text_id = self
            .0
            .service()
            .media()
            .text()
            .save(&request.user_id, request.text_id, &request.text)
            .await
            .map_err(|e| HandlerError::from(e.to_string()))?;

        Ok(Json(SaveTextResponse {
            text_id,
            context_key: ctx.key(),
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::MockAppState;
    use twinqit_api_domain::{
        MockService,
        media::{
            TextId,
            service::{MockMediaService, text::MockTextService},
        },
        user::UserId,
    };
    use twinqit_api_util::restate::MockObjectContextProvider;

    #[tokio::test]
    async fn test_save_success() {
        let mut state_mock = MockAppState::new();
        let mut service_mock = MockService::new();
        let mut media_service_mock = MockMediaService::new();
        let mut text_service_mock = MockTextService::new();
        let mut ctx_mock = MockObjectContextProvider::new();

        let user_id = UserId::default();
        let text_id = Some(TextId::default());
        let text = "Hello, world!";
        let expected_text_id = text_id.clone();

        text_service_mock
            .expect_save()
            .with(
                mockall::predicate::eq(user_id.clone()),
                mockall::predicate::eq(text_id.clone()),
                mockall::predicate::eq(text),
            )
            .return_once(move |_, _, _| Ok(expected_text_id.clone()));

        media_service_mock
            .expect_text()
            .return_once(move || Arc::new(text_service_mock));

        service_mock
            .expect_media()
            .return_once(move || Arc::new(media_service_mock));

        state_mock.expect_service().return_once(move || Arc::new(service_mock));

        ctx_mock.expect_key().return_once(|| "test-key".to_string());

        let handler = MediaTextObjectImpl(Arc::new(state_mock));
        let request = SaveTextRequest {
            user_id,
            text_id: text_id.clone(),
            text: text.to_string(),
        };

        let result = handler.save_internal(&ctx_mock, Json(request)).await;

        assert!(result.is_ok());

        let response = result.unwrap().0;
        assert_eq!(response.text_id, text_id);
        assert_eq!(response.context_key, "test-key");
    }

    #[tokio::test]
    async fn test_save_service_error() {
        let mut state_mock = MockAppState::new();
        let mut service_mock = MockService::new();
        let mut media_service_mock = MockMediaService::new();
        let mut text_service_mock = MockTextService::new();
        let mut ctx_mock = MockObjectContextProvider::new();

        let user_id = UserId::default();
        let text_id = Some(TextId::default());
        let text = "Hello, world!";

        text_service_mock
            .expect_save()
            .return_once(|_, _, _| Err(anyhow::anyhow!("Database error")));

        media_service_mock
            .expect_text()
            .return_once(move || Arc::new(text_service_mock));

        service_mock
            .expect_media()
            .return_once(move || Arc::new(media_service_mock));

        state_mock.expect_service().return_once(move || Arc::new(service_mock));

        ctx_mock.expect_key().return_once(|| "test-key".to_string());

        let handler = MediaTextObjectImpl(Arc::new(state_mock));
        let request = SaveTextRequest {
            user_id,
            text_id,
            text: text.to_string(),
        };

        let result = handler.save_internal(&ctx_mock, Json(request)).await;

        assert!(result.is_err());
        assert!(matches!(result, Err(HandlerError { .. })));
    }
}

But you can only do if for methods that are defined for the ObjectContext impl:

impl<'ctx> ObjectContext<'ctx> {
    /// Get object key.
    pub fn key(&self) -> &str {
        &self.key
    }

    /// Get request headers.
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }

    /// Get request headers.
    pub fn headers_mut(&mut self) -> &HeaderMap {
        &mut self.headers
    }
}

From a design perspective, if it were possible to define the other methods directly on the context, users could provide context providers that they could use for mocking.

Even if this were not possible from a design perspective, I believe it would still be feasible to include context providers in the SDK, as these providers would have access to ContextInnernal:

pub struct ObjectContext<'ctx> {
    key: String,
    random_seed: u64,
    #[cfg(feature = "rand")]
    std_rng: rand::prelude::StdRng,
    headers: HeaderMap,
    pub(crate) inner: &'ctx ContextInternal,
}

Then you could do something like:

#[automock]
pub trait ObjectContextProvider {
    fn key(&self) -> String;
    fn rand_uuid(&mut self) -> uuid::Uuid;
}

impl ObjectContextProvider for ObjectContext<'_> {
    fn key(&self) -> String {
        self.key().to_owned()
    }
    fn rand_uuid(&mut self) -> uuid::Uuid {
        self.inner.rand_uuid()
    }
}

However, this had the disadvantage that users were forced to use a mocking library that you had to decide on.

[BrynCooke]

I initially tried to implement my own versions of things like ContextClient, but these traits are currently sealed and therefore cannot be implemented.

pub trait ContextClient<'ctx>: private::SealedContext<'ctx> {

[slinkydeveloper]

I initially tried to implement my own versions of things like ContextClient, but these traits are currently sealed and therefore cannot be implemented.

Ah that's right I completely forgot about that. The thing is, it makes sense for those traits to be sealed, but i can also see the point about mocking... I need to think this through, in combination with the changes to the macro...