From 78e568db247f96b7627b6709bd5c9710e9f680ef Mon Sep 17 00:00:00 2001
From: Joe Hoyle <joehoyle@gmail.com>
Date: Tue, 5 Mar 2024 16:50:46 +0200
Subject: [PATCH] Fix tests

---
 tests/test-lambda.ts | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/test-lambda.ts b/tests/test-lambda.ts
index adaf7a3..73fa969 100644
--- a/tests/test-lambda.ts
+++ b/tests/test-lambda.ts
@@ -32,6 +32,7 @@ class TestResponseStream {
 	contentType: string | undefined;
 	body: string | Buffer | undefined;
 	headers: { [key: string]: string } = {};
+	metadata: any;
 
 	setContentType( type: string ): void {
 		this.contentType = type;
@@ -46,7 +47,9 @@ class TestResponseStream {
 		}
 	}
 	end(): void {
-
+		if ( this.metadata.headers['Content-Type'] ) {
+			this.contentType = this.metadata.headers['Content-Type'];
+		}
 	}
 }
 
@@ -63,7 +66,8 @@ global.awslambda = {
 		 * @param stream The response stream.
 		 * @param metadata The metadata for the response.
 		 */
-		from( stream: ResponseStream, metadata ) : ResponseStream {
+		from( stream: TestResponseStream, metadata ) : TestResponseStream {
+			stream.metadata = metadata;
 			return stream;
 		},
 	},