Support Lambda Streaming (#229)

* feat: first impl of streaming

* feat: support usecase without awslambda.HttpResponseStream.from

* fix: swap performance.now for Date.now

* fix: there's another one

* fix: only assign globalThis if it doesn't exist yet

* add docstring

Author: Skn0tt

src/lambdalocal.ts

@@ -13,6 +13,7 @@ import os = require('os');
 import { createServer, IncomingMessage, ServerResponse } from 'http';
 import utils = require('./lib/utils.js');
 import Context = require('./lib/context.js');
+require("./lib/streaming.js");
 
 /*
  * Lambda local version

src/lib/streaming.ts

@@ -0,0 +1,64 @@
+/**
+ * Implements Lambda Response Streaming by polyfilling
+ * `awslambda.streamifyResponse` and `awslambda.HttpResponseStream.from`.
+ * 
+ * If they're used, `execute` will return a `ReadableStream` as `body`.
+ * 
+ * See https://aws.amazon.com/fr/blogs/compute/introducing-aws-lambda-response-streaming/ for reference.
+ */
+
+import { PassThrough } from "stream";
+
+function streamifyResponse(handler) {
+  return (event, context) =>
+    new Promise(async (resolve, reject) => {
+      const body = new StreamingBody(resolve);
+
+      try {
+        const metadata = await handler(event, body, context); // cb not supported
+        if (!body.headersSent) {
+          body.sendHeader(metadata)
+        }
+      } catch (error) {
+        reject(error);
+      }
+    });
+}
+
+class StreamingBody extends PassThrough {
+  constructor(private readonly resolve: (metadata) => void) {
+    super();
+  }
+
+  public headersSent = false
+  sendHeader(metadata: any = {}) {
+    const headers = { ...metadata.headers }
+    if (this.contentType) {
+      headers["Content-Type"] = this.contentType
+    }
+
+    this.resolve({
+      ...metadata,
+      headers,
+      body: this
+    })
+    this.headersSent = true
+  }
+
+  private contentType: string;
+  setContentType(contentType) {
+    this.contentType = contentType;
+  }
+}
+
+class HttpResponseStream {
+  static from(responseStream: StreamingBody, metadata) {
+    responseStream.sendHeader(metadata);
+    return responseStream;
+  }
+}
+
+globalThis.awslambda = globalThis.awslambda || {
+  streamifyResponse,
+  HttpResponseStream,
+};

test/functs/test-func-streaming-simple.js

@@ -0,0 +1,15 @@
+/*
+ * Lambda function used for streaming test.
+ */
+exports.handler = awslambda.streamifyResponse(
+    async (event, responseStream, context) => {
+        responseStream.setContentType("text/plain");
+        
+        responseStream.write("foo");
+        setTimeout(() => {
+            responseStream.write("bar");
+            responseStream.end();
+        }, 100);
+    }
+);
+

test/functs/test-func-streaming.js

@@ -0,0 +1,23 @@
+/*
+ * Lambda function used for streaming test.
+ */
+exports.handler = awslambda.streamifyResponse(
+    async (event, responseStream, context) => {
+        const metadata = {
+            statusCode: 200,
+            headers: {
+                "X-Foo": "Bar"
+            }
+        };
+
+        responseStream.setContentType("text/plain");
+        responseStream = awslambda.HttpResponseStream.from(responseStream, metadata);
+        
+        responseStream.write("foo");
+        setTimeout(() => {
+            responseStream.write("bar");
+            responseStream.end();
+        }, 100);
+    }
+);
+

test/test.js

@@ -420,6 +420,67 @@ describe("- Testing lambdalocal.js", function () {
                 });
             });
         });
+
+        describe('* Streaming Response', function () {
+            it('should return a readable stream as `body`', function () {
+                var lambdalocal = require(lambdalocal_path);
+                lambdalocal.setLogger(winston);
+                return lambdalocal.execute({
+                    event: require(path.join(__dirname, "./events/test-event.js")),
+                    lambdaPath: path.join(__dirname, "./functs/test-func-streaming.js"),
+                    lambdaHandler: functionName,
+                    callbackWaitsForEmptyEventLoop: false,
+                    timeoutMs: timeoutMs,
+                    verboseLevel: 1
+                }).then(function (data) {
+                    assert.deepEqual(
+                        data.headers,
+                        { "Content-Type": "text/plain", "X-Foo": "Bar" }
+                    );
+
+                    return new Promise((resolve, reject) => {
+                        const chunks = []
+                        const times = []
+                        data.body.on('data', (chunk) => {
+                            chunks.push(chunk.toString())
+                            times.push(Date.now())
+                        });
+                        data.body.on("end", () => {
+                            assert.deepEqual(chunks, ["foo", "bar"])
+                            assert.closeTo(times[1] - times[0], 100, 50)
+                            resolve()
+                        });
+                    })
+                })
+            });
+
+            it('also works without calling HttpResponseStream.from', function () {
+                var lambdalocal = require(lambdalocal_path);
+                lambdalocal.setLogger(winston);
+                return lambdalocal.execute({
+                    event: require(path.join(__dirname, "./events/test-event.js")),
+                    lambdaPath: path.join(__dirname, "./functs/test-func-streaming-simple.js"),
+                    lambdaHandler: functionName,
+                    callbackWaitsForEmptyEventLoop: false,
+                    timeoutMs: timeoutMs,
+                    verboseLevel: 1
+                }).then(function (data) {
+                    return new Promise((resolve, reject) => {
+                        const chunks = []
+                        const times = []
+                        data.body.on('data', (chunk) => {
+                            chunks.push(chunk.toString())
+                            times.push(Date.now())
+                        });
+                        data.body.on("end", () => {
+                            assert.deepEqual(chunks, ["foo", "bar"])
+                            assert.closeTo(times[1] - times[0], 100, 50)
+                            resolve()
+                        });
+                    })
+                })
+            });
+        });
     }
 });
 describe("- Testing cli.js", function () {