Also mock the content type from headers

humanmade · Mar 5, 2024 · 6dda3aa · 6dda3aa
1 parent 78e568d
commit 6dda3aa
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/global.d.ts b/global.d.ts
@@ -2,6 +2,7 @@ declare type ResponseStream = {
 	setContentType( type: string ): void;
 	write( stream: string | Buffer ): void;
 	end(): void;
+	metadata?: any;
 };
 
 declare type StreamifyHandler = ( event: APIGatewayProxyEventV2, response: ResponseStream ) => Promise<any>;

diff --git a/src/lambda-handler.ts b/src/lambda-handler.ts
@@ -97,6 +97,7 @@ if ( typeof awslambda === 'undefined' ) {
 			from( response: ResponseStream, metadata: {
 				headers?: Record<string, string>,
 			} ): ResponseStream {
+				response.metadata = metadata;
 				return response;
 			},
 		},

diff --git a/tests/test-private-upload.ts b/tests/test-private-upload.ts
@@ -74,6 +74,9 @@ test( 'Test get private upload', async () => {
 		 * End the response.
 		 */
 		end(): void {
+			if ( this.metadata.headers['Cache-Control'] ) {
+				contentType = this.metadata.headers['Content-Type'];
+			}
 		},
 	} );
 
@@ -115,6 +118,9 @@ test( 'Test get private upload with presign params', async () => {
 		 * End the response.
 		 */
 		end(): void {
+			if ( this.metadata.headers['Cache-Control'] ) {
+				contentType = this.metadata.headers['Content-Type'];
+			}
 		},
 	} );