Merge pull request #4 from cloudnc/feat/request-context-for-response

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> {
@@ -38,10 +34,15 @@ function unframeRequest(requestBody: Uint8Array): Uint8Array {
   return new Uint8Array(requestBody).slice(5);
 }
 
+export function readGrpcRequest(request: Request): Uint8Array | null {
+  const requestBody = request.postDataBuffer();
+  return !requestBody ? null : unframeRequest(requestBody);
+}
+
 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 +53,14 @@ 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);
+    const grpcResponse = typeof response === 'function'
+      ? response(readGrpcRequest(route.request()))
+      : response;
+
+    const grpcResponseBody = grpcResponseToBuffer(grpcResponse);
 
     return route.fulfill({
-      body: grpcResponse,
+      body: grpcResponseBody,
       contentType: 'application/grpc-web+proto',
       headers: {
         'Access-Control-Allow-Origin': '*',
@@ -72,8 +76,7 @@ export function mockGrpcUnary(
         }
 
         if (requestPredicate) {
-          const messageBody = req.postDataBuffer();
-          const unframed = !messageBody ? null : unframeRequest(messageBody);
+          const unframed = readGrpcRequest(req);
           return requestPredicate(unframed, req);
         }
 
@@ -82,8 +85,7 @@ export function mockGrpcUnary(
 
       await page.waitForResponse((resp) => resp.url().includes(url));
 
-      const messageBody = request.postDataBuffer();
-      const requestMessage = !messageBody ? null : unframeRequest(messageBody);
+      const requestMessage = readGrpcRequest(request);
 
       return { requestMessage };
     },