test(eventstream-handler-node): mock EventSigningStream in beforeEach

Author: trivikr

packages/eventstream-handler-node/src/EventStreamPayloadHandler.spec.ts

@@ -1,14 +1,12 @@
+import { Decoder, Encoder, EventSigner, FinalizeHandler, FinalizeHandlerArguments, HttpRequest } from "@aws-sdk/types";
 import { once } from "events";
 import { PassThrough, Readable } from "stream";
 
-const mockSingingStream = jest.fn().mockImplementation(() => new PassThrough());
-jest.mock("./EventSigningStream", () => ({
-  EventSigningStream: mockSingingStream,
-}));
-import { Decoder, Encoder, EventSigner, FinalizeHandler, FinalizeHandlerArguments, HttpRequest } from "@aws-sdk/types";
-
+import { EventSigningStream } from "./EventSigningStream";
 import { EventStreamPayloadHandler } from "./EventStreamPayloadHandler";
 
+jest.mock("./EventSigningStream");
+
 describe("EventStreamPayloadHandler", () => {
   const mockSigner: EventSigner = {
     sign: jest.fn(),
@@ -18,6 +16,7 @@ describe("EventStreamPayloadHandler", () => {
   const mockNextHandler: FinalizeHandler<any, any> = jest.fn();
 
   beforeEach(() => {
+    (EventSigningStream as unknown as jest.Mock).mockImplementation(() => new PassThrough());
     jest.clearAllMocks();
   });
 
@@ -36,24 +35,24 @@ describe("EventStreamPayloadHandler", () => {
   });
 
   it("should close the request payload if downstream middleware throws", async () => {
-    expect.assertions(2);
-    (mockNextHandler as any).mockImplementationOnce(() => Promise.reject(new Error()));
+    const mockError = new Error("mockError");
+    (mockNextHandler as any).mockImplementationOnce(() => Promise.reject(mockError));
     const handler = new EventStreamPayloadHandler({
       eventSigner: () => Promise.resolve(mockSigner),
       utf8Decoder: mockUtf8Decoder,
       utf8Encoder: mockUtf8encoder,
     });
     const mockRequest = { body: new Readable() } as HttpRequest;
-    let error;
+
     try {
       await handler.handle(mockNextHandler, {
         request: mockRequest,
         input: {},
       });
-    } catch (e) {
-      error = e;
+      fail(`Expected ${mockError} to be thrown.`);
+    } catch (error) {
+      expect(error).toBe(mockError);
     }
-    expect(error instanceof Error).toBe(true);
 
     // Expect stream is closed
     // Ref: should use writableEnded when bumped to Node 13+
@@ -66,23 +65,31 @@ describe("EventStreamPayloadHandler", () => {
   });
 
   it("should call event signer with request signature from signing middleware", async () => {
-    const authorization =
-      "AWS4-HMAC-SHA256 Credential=AKID/20200510/us-west-2/foo/aws4_request, SignedHeaders=host, Signature=1234567890";
+    const priorSignature = "1234567890";
+    const authorization = `AWS4-HMAC-SHA256 Credential=AKID/20200510/us-west-2/foo/aws4_request, SignedHeaders=host, Signature=${priorSignature}`;
+
     const mockRequest = {
       body: new PassThrough(),
       headers: { authorization },
     } as any;
+
     const handler = new EventStreamPayloadHandler({
       eventSigner: () => Promise.resolve(mockSigner),
       utf8Decoder: mockUtf8Decoder,
       utf8Encoder: mockUtf8encoder,
     });
+
     await handler.handle(mockNextHandler, {
       request: mockRequest,
       input: {},
     });
-    expect(mockSingingStream.mock.calls.length).toBe(1);
-    expect(mockSingingStream.mock.calls[0][0].priorSignature).toBe("1234567890");
+
+    expect(EventSigningStream).toHaveBeenCalledTimes(1);
+    expect(EventSigningStream).toHaveBeenCalledWith({
+      priorSignature,
+      eventMarshaller: expect.anything(),
+      eventSigner: expect.anything(),
+    });
   });
 
   it("should start piping to request payload through event signer if downstream middleware returns", async () => {

packages/eventstream-handler-node/src/EventStreamPayloadHandler.ts

@@ -32,6 +32,7 @@ export interface EventStreamPayloadHandlerOptions {
 export class EventStreamPayloadHandler implements IEventStreamPayloadHandler {
   private readonly eventSigner: Provider<EventSigner>;
   private readonly eventMarshaller: EventMarshaller;
+
   constructor(options: EventStreamPayloadHandlerOptions) {
     this.eventSigner = options.eventSigner;
     this.eventMarshaller = new EventMarshaller(options.utf8Encoder, options.utf8Decoder);
@@ -45,13 +46,16 @@ export class EventStreamPayloadHandler implements IEventStreamPayloadHandler {
   ): Promise<FinalizeHandlerOutput<T>> {
     const request = args.request as HttpRequest;
     const { body: payload } = request;
+
     if (!(payload instanceof Readable)) {
       throw new Error("Eventstream payload must be a Readable stream.");
     }
+
     const payloadStream = payload as Readable;
     request.body = new PassThrough({
       objectMode: true,
     });
+
     let result: FinalizeHandlerOutput<any>;
     try {
       result = await next(args);
@@ -61,6 +65,7 @@ export class EventStreamPayloadHandler implements IEventStreamPayloadHandler {
       request.body.end();
       throw e;
     }
+
     // If response is successful, start piping the payload stream
     const match = (request.headers["authorization"] || "").match(/Signature=([\w]+)$/);
     // Sign the eventstream based on the signature from initial request.
@@ -70,11 +75,13 @@ export class EventStreamPayloadHandler implements IEventStreamPayloadHandler {
       eventMarshaller: this.eventMarshaller,
       eventSigner: await this.eventSigner(),
     });
+
     pipeline(payloadStream, signingStream, request.body, (err) => {
       if (err) {
         throw err;
       }
     });
+
     return result;
   }
 }