Add support for returning trailers in thenCallback handlers

Author: pimterry

src/rules/requests/request-handler-definitions.ts

@@ -7,6 +7,7 @@ import { stripIndent } from 'common-tags';
 
 import {
     Headers,
+    Trailers,
     CompletedRequest,
     CompletedBody,
     Explainable,
@@ -180,6 +181,15 @@ export interface CallbackResponseMessageResult {
      */
     headers?: Headers;
 
+    /**
+     * The replacement HTTP trailers, as an object of string keys and either
+     * single string or array of string values. Note that there are not all
+     * header fields are valid as trailers, and there are other requirements
+     * such as chunked encoding that must be met for trailers to be sent
+     * successfully.
+     */
+    trailers?: Trailers;
+
     /**
      * A string or buffer, which replaces the response body if set. This will
      * be automatically encoded to match the Content-Encoding defined in your

src/rules/requests/request-handlers.ts

@@ -166,9 +166,14 @@ export class SimpleHandler extends SimpleHandlerDefinition {
     }
 }
 
-async function writeResponseFromCallback(result: CallbackResponseMessageResult, response: OngoingResponse) {
+async function writeResponseFromCallback(
+    result: CallbackResponseMessageResult,
+    response: OngoingResponse
+) {
     if (result.json !== undefined) {
-        result.headers = _.assign(result.headers || {}, { 'Content-Type': 'application/json' });
+        result.headers = Object.assign(result.headers || {}, {
+            'Content-Type': 'application/json'
+        });
         result.body = JSON.stringify(result.json);
         delete result.json;
     }
@@ -193,6 +198,9 @@ async function writeResponseFromCallback(result: CallbackResponseMessageResult,
         result.statusMessage,
         result.headers
     );
+
+    if (result.trailers) response.addTrailers(result.trailers);
+
     response.end(result.rawBody || "");
 }
 

src/types.ts

@@ -36,6 +36,11 @@ export interface Headers {
     [key: string]: undefined | string | string[];
 }
 
+export interface Trailers {
+    // 0+ of any trailer
+    [key: string]: undefined | string | string[];
+}
+
 export type RawHeaders = Array<[key: string, value: string]>;
 
 export interface Request {

test/integration/handlers/callback-response.spec.ts

@@ -1,6 +1,15 @@
 import * as zlib from 'zlib';
+import * as http from 'http';
+
 import { getLocal } from "../../..";
-import { expect, fetch, isNode, isWeb, headersToObject } from "../../test-utils";
+import {
+    expect,
+    fetch,
+    isNode,
+    isWeb,
+    headersToObject,
+    nodeOnly
+} from "../../test-utils";
 
 describe("Callback response handlers", function () {
 
@@ -71,6 +80,34 @@ describe("Callback response handlers", function () {
         expect(await response.text()).to.equal("Error: Cannot set custom :status pseudoheader values");
     });
 
+    nodeOnly(() => { // Browsers can't read trailers
+        it("should allow mocking response trailers with a callback", async () => {
+            await server.forGet("/mocked-endpoint").thenCallback(() => {
+                return {
+                    statusCode: 200,
+                    trailers: {
+                        'mock-trailer': 'set'
+                    }
+                };
+            });
+
+            const response = await new Promise<http.IncomingMessage>((resolve, reject) => {
+                const req = http.request(server.urlFor("/mocked-endpoint")).end();
+                req.on('response', (res) => {
+                    // Wait until everything is 100% done, so we definitely have trailers
+                    res.resume();
+                    res.on('end', () => resolve(res));
+                });
+                req.on('error', reject);
+            });
+
+            expect(response.statusCode).to.equal(200);
+            expect(response.trailers).to.deep.equal({
+                'mock-trailer': 'set'
+            });
+        });
+    });
+
     it("should allow mocking body as json with callback", async () => {
         await server.forGet("/mocked-endpoint").thenCallback(() => ({
             statusCode: 201,