feat(mockGrpcUnary): Add functionality to make the mocked response affected by the request data

Author: zak-cloudnc

src/base/index.ts

@@ -1,4 +1,4 @@
-import { status as Status } from 'grpc';
+import {status as Status} from 'grpc';
 
 export interface GrpcErrorResponse {
   status: Status;
@@ -9,6 +9,8 @@ export interface GrpcSuccessResponse {
   message: Uint8Array;
 }
 
+export type GrpcResponse = GrpcSuccessResponse | GrpcErrorResponse
+
 function fourBytesLength(sized: { length: number }): Uint8Array {
   const arr = new Uint8Array(4); // an Int32 takes 4 bytes
   const view = new DataView(arr.buffer);
@@ -23,8 +25,10 @@ export class GrpcUnknownStatus extends Error {
 }
 
 export function grpcResponseToBuffer(
-  response: GrpcSuccessResponse | GrpcErrorResponse
+  response: GrpcResponse
 ): Buffer {
+
+  // error messages need to have a zero length message field to be considered valid
   const message = 'message' in response ? response.message : new Uint8Array();
 
   // all success responses have status OK

src/playwright/index.ts

@@ -1,11 +1,7 @@
-import { expect, Page } from '@playwright/test';
-import { grpc } from '@improbable-eng/grpc-web';
-import {
-  GrpcErrorResponse,
-  grpcResponseToBuffer,
-  GrpcSuccessResponse,
-} from '../base';
-import { Request } from 'playwright-core';
+import {expect, Page} from '@playwright/test';
+import {grpc} from '@improbable-eng/grpc-web';
+import {GrpcResponse, grpcResponseToBuffer,} from '../base';
+import {Request} from 'playwright-core';
 
 export interface UnaryMethodDefinitionish
   extends grpc.UnaryMethodDefinition<any, any> {
@@ -41,7 +37,7 @@ function unframeRequest(requestBody: Uint8Array): Uint8Array {
 export function mockGrpcUnary(
   page: Page,
   rpc: UnaryMethodDefinitionish,
-  response: GrpcSuccessResponse | GrpcErrorResponse
+  response: GrpcResponse | ((request: Uint8Array|null) => GrpcResponse)
 ): MockedGrpcCall {
   const url = `/${rpc.service.serviceName}/${rpc.methodName}`;
 
@@ -52,11 +48,19 @@ export function mockGrpcUnary(
       'ALL gRPC requests should be a POST request'
     ).toBe('POST');
 
-    // error messages need to have a zero length message field to be considered valid
-    const grpcResponse = grpcResponseToBuffer(response);
+    let grpcResponse: GrpcResponse
+
+    if (typeof response === 'function') {
+      const requestBody = route.request().postDataBuffer();
+      grpcResponse = response(!requestBody ? null : unframeRequest(requestBody));
+    } else {
+      grpcResponse = response;
+    }
+
+    const grpcResponseBody = grpcResponseToBuffer(grpcResponse)
 
     return route.fulfill({
-      body: grpcResponse,
+      body: grpcResponseBody,
       contentType: 'application/grpc-web+proto',
       headers: {
         'Access-Control-Allow-Origin': '*',