Explicitly Route Messages with Wolverine

TL;DR: Wolverine message handler signatures can lead to easier unit testing code than comparable “IHandler of T” frameworks.

Most of the time I think you can just allow Wolverine to handle message routing for you with some simple configured rules or conventions. However, once in awhile you’ll need to override those rules and tell Wolverine exactly where those messages should go.

I’ve been working (playing) with Midjourney quite a bit lately trying to make images for the JasperFx Software website. You can try it out to generate images for free, but the image generation gets a lower priority than their paying customers. Using that as an example, let’s say that we were using Wolverine to build our own Midjourney clone. At some point, there’s maybe an asynchronous message handler like this one that takes a request to generate a new image based on the user’s prompt, but routes the actual work to either a higher or lower priority queue based on whether the user is a premium customer:

    public record GenerateImage(string Prompt, Guid ImageId);

    public record ImageRequest(string Prompt, string CustomerId);

    public record ImageGenerated(Guid Id, byte[] Image);

    public class Customer
    {
        public string Id { get; set; }
        public bool PremiumMembership { get; set; }
    }

    public class ImageSaga : Saga
    {
        public Guid Id { get; set; }
        
        public string CustomerId { get; set; }

        public Task Handle(ImageGenerated generated)
        {
            // look up the customer, figure out how to send the
            // image to their client.
            throw new NotImplementedException("Not done yet:)");
            
            MarkCompleted();
        }
    }
    
    public static class GenerateImageHandler
    {
        // I'm assuming the usage of Marten middleware here
        // to handle transactions and the outbox mechanics
        public static async Task HandleAsync(
            ImageRequest request, 
            IDocumentSession session, 
            IMessageBus messageBus,
            CancellationToken cancellationToken)
        {
            var customer = await session
                .LoadAsync<Customer>(request.CustomerId, cancellationToken);

            // I'm starting a new saga to track the state of the 
            // image when we get the callback from the downstream
            // image generation service
            var imageSaga = new ImageSaga();
            session.Insert(imageSaga);

            var outgoing = new GenerateImage(request.Prompt, imageSaga.Id);
            if (customer.PremiumMembership)
            {
                // Send the message to a named endpoint we've configured for the faster
                // processing
                await messageBus.EndpointFor("premium-processing")
                    .SendAsync(outgoing);
            }
            else
            {
                // Send the message to a named endpoint we've configured for slower
                // processing
                await messageBus.EndpointFor("basic-processing")
                    .SendAsync(outgoing);
            }
        }
    }

A couple notes on the code above:

  • I’m assuming the usage of Marten for persistence (of course), with the auto transactional middleware policy applied
  • I’ve configured a PostgreSQL backed outbox for Wolverine
  • It’s likely a slow process, so I’m assuming there’s going to be an asynchronous callback from the actual image generator later. I’m leveraging Wolverine’s stateful saga support to track the customer of the original image for processing later

Wolverine V1.3 dropped today with a little improvement for exactly this scenario (based on some usage by a JasperFx client) so you can use cascading messages instead of having to deal directly with the IMessageBus service. Let’s rewrite the explicit code up above, but this time try to turn the actual routing logic into a pure function that could be easy to unit test:

    public static class GenerateImageHandler
    {
        // Using Wolverine's compound handlers to remove all the asynchronous
        // junk from the main Handle() method
        public static Task<Customer> LoadAsync(
            ImageRequest request, 
            IDocumentSession session,
            CancellationToken cancellationToken)
        {
            return session.LoadAsync<Customer>(request.CustomerId, cancellationToken);
        }
        
        
        public static (RoutedToEndpointMessage<GenerateImage>, ImageSaga) Handle(
            ImageRequest request, 
            Customer customer)
        {

            // I'm starting a new saga to track the state of the 
            // image when we get the callback from the downstream
            // image generation service
            var imageSaga = new ImageSaga
            {
                // I need to assign the image id in memory
                // to make this all work
                Id = CombGuidIdGeneration.NewGuid()
            };

            var outgoing = new GenerateImage(request.Prompt, imageSaga.Id);
            var destination = customer.PremiumMembership ? "premium-processing" : "basic-processing";
            
            return (outgoing.ToEndpoint(destination), imageSaga);
        }
    }

The handler above is the equivalent in functionality to the earlier version. It’s not really that much less code, but I think it’s a bit more declarative. What’s most important to me is the potential for unit testing the decision about where the customer requests go as shown in this fake test:

    [Fact]
    public void should_send_the_request_to_premium_processing_for_premium_customers()
    {
        var request = new ImageRequest("a wolverine ice skating in the country side", "alice");
        var customer = new Customer
        {
            Id = "alice",
            PremiumMembership = true
        };

        var (command, image) = GenerateImageHandler.Handle(request, customer);
        
        command.EndpointName.ShouldBe("premium-processing");
        command.Message.Prompt.ShouldBe(request.Prompt);
        command.Message.ImageId.ShouldBe(image.Id);
        
        image.CustomerId.ShouldBe(request.CustomerId);
    }

What I’m hoping you take away from that code sample is that testing the logic part of the ImageRequest message processing turns into a simple state-based test — meaning that you’re just pushing in the known inputs and measuring the values returned by the method. You’d still need to pair this unit test with a full integration test, but at least you’d know that the routing logic is correct before you wrestle with potential integration issues.

Leave a comment