Mock connect-es services for RTK Query in Storybook

[boan-anbo]

I feel this is a common stack for React apps so I'm sharing my current implementation for mocking services generated by connect-es here for those who need it.

This showcases both how to use RTQ Query with connect-es grpc-web clients and how to mock both in Storybook environment.

// Utility function to create either a real or mock client based on the environment.
function createClient(service: any, transport = grpcWebTransport): PromiseClient<any> {
    // Check if the environment is Storybook
    if (getDevEnv('STORYBOOK_ENV') === 'storybook') {
        console.warn("Mocking client for service: " + service.typeName);
        // Create a mock transport using `createRouterTransport` from `@bufbuild/connect`
        const mockTransport = createRouterTransport(({service: mockService}) => {
            // For each service, define the methods using `mockService`. If a mock implementation
            // is provided in `mockImplementations`, use that. Otherwise, fallback to the real implementation.
            mockService(service, {
                ...Object.fromEntries(
                    // Get the method names from the `service.methods` object. This assumes `service.methods`
                    // is an object where keys are method names.
                    Object.keys(service.methods || {}).map(methodName => [
                        methodName,
                        // Define a function for each method name.
                        // @ts-ignore
                        (...args) => {
                            // Check if a mock implementation is provided for this method
                            // @ts-ignore
                            const mockImpl = mockImplementations[service.typeName]?.[methodName];
                            if (mockImpl) {
                                // If a mock implementation is provided, use that.
                                return mockImpl(...args);
                            } else {
                                // If no mock implementation is provided, log a warning and use the real implementation.
                                console.warn(`No mock implementation for ${methodName} under ${service.typeName}, using real implementation.`);
                                // @ts-ignore
                                // Call the real implementation of the method on the real client.
                                return createPromiseClient(service, transport)[methodName](...args);
                            }
                        }
                    ])
                ),
                // If there are any specific mock implementations provided in `mockImplementations`,
                // use those to override the generated functions above.
                // @ts-ignore
                ...mockImplementations[service.typeName]
            });
        });
        return createPromiseClient(service, mockTransport);
    }
    // If not in Storybook, create and return a real client
    return createPromiseClient(service, transport);
}

// Object containing mock implementations for different services
const mockImplementations = {
    [WorkflowFireflyService.typeName]: {
        // Mock implementation for the `deleteFireflies` method
        deleteFireflies: () => new DeleteFirefliesResponse({
            affectedCount: 1,
            success: true
        }),
        getFireflyById: () => new GetFireflyByIdResponse({
            firefly: new Firefly({
                uuid: "123",
            })
        }),
        listFireflies: () => new ListFirefliesResponse({
            searchResults: getSampleFireflyResult(10)
        })
    },
    // ... other services
};


// Client for WorkflowFireflyService
export const workflowFireflyClient = createClient(WorkflowFireflyService) as PromiseClient<typeof WorkflowFireflyService>;

The connect-es generated clients will then be wrapped in RTK Query as custom queryFn like so:

export const carrelApi = createApi({
    reducerPath: "carrelApi",
    baseQuery: fetchBaseQuery({baseUrl: ''}),
    endpoints: (builder) => ({
        workflowFirefliesDeleteFireflies: builder.mutation<DeleteFirefliesResponse, DeleteFirefliesRequest>({
            queryFn: async (args) => {
                return await clientCallWrapper(workflowFireflyClient.deleteFireflies, args);
            },
        }),
        workflowFirefliesListFireflies: builder.query<ListFirefliesResponse, ListFirefliesRequest>({
                queryFn: async (args) => {
                    return await clientCallWrapper(workflowFireflyClient.listFireflies, args);
                }
            }
        ),
        workflowFirefliesSearchFireflies: builder.query<SearchFirefliesResponse, SearchFirefliesRequest>({
                queryFn: async (args) => {
                    return await clientCallWrapper(workflowFireflyClient.searchFireflies, args);
                }
            }
        ),
        workflowFirefliesExtractPDFFireflies: builder.query<ExtractPdfFirefliesResponse, ExtractPdfFirefliesRequest>({
            queryFn: async (args) => {
                return await clientCallWrapper(workflowFireflyClient.extractPdfFireflies, args);
            }
        }),
        workflowFirefliesGetFireflyById: builder.query<GetFireflyByIdResponse, GetFireflyByIdRequest>({
            queryFn: async (args) => {
                return await clientCallWrapper(workflowFireflyClient.getFireflyById, args);

            }
        }),
    })
})


async function clientCallWrapper<T, U>(client: (request: T) => Promise<U>, request: T): Promise<{
    data: U
} | ReturnType<typeof handleClientError>> {
    try {
        let response = await client(request);
        return {data: response}
    } catch (e) {
        return handleClientError(e);
    }
}

// And finally a handler which converts `ConnectError` to the expected error type for RTQ Query
function handleClientError(e: unknown) {
    if (e instanceof ConnectError) {
        return {
            error: {
                status: e.code,
                statusText: e.message,
                data: e.details,
            },
        }
    } else {
        return {
            error: {
                status: 500,
                statusText: 'Unknown RTK Query Error',
                data: e,
            },
        }
    }
}