feat(rest): allow controller methods to handle response writing

loopbackio · Sep 25, 2018 · 2bfd50e · 2bfd50e
1 parent 7e1acfc
commit 2bfd50e
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 2 deletions.
diff --git a/packages/rest/src/writer.ts b/packages/rest/src/writer.ts
@@ -19,13 +19,21 @@ export function writeResultToResponse(
   // result returned back from invoking controller method
   result: OperationRetval,
 ): void {
-  if (!result) {
+  // Bypass response writing if the controller method returns `response` itself
+  // or the response headers have been sent
+  if (result === response || response.headersSent) {
+    return;
+  }
+  if (result === undefined) {
     response.statusCode = 204;
     response.end();
     return;
   }
 
-  if (result instanceof Readable || typeof result.pipe === 'function') {
+  const isStream =
+    result instanceof Readable || typeof (result && result.pipe) === 'function';
+
+  if (isStream) {
     response.setHeader('Content-Type', 'application/octet-stream');
     // Stream
     result.pipe(response);

diff --git a/packages/rest/test/unit/writer.unit.ts b/packages/rest/test/unit/writer.unit.ts
@@ -68,6 +68,13 @@ describe('writer', () => {
     expect(result.payload).to.equal('ABC123');
   });
 
+  it('writes null object result to response as JSON', async () => {
+    writeResultToResponse(response, null);
+    const result = await observedResponse;
+    expect(result.headers['content-type']).to.eql('application/json');
+    expect(result.payload).to.equal('null');
+  });
+
   it('sends 204 No Content when the response is undefined', async () => {
     writeResultToResponse(response, undefined);
     const result = await observedResponse;
@@ -77,6 +84,36 @@ describe('writer', () => {
     expect(result.payload).to.equal('');
   });
 
+  it('skips writing when the return value is the response', async () => {
+    response
+      .status(200)
+      .contentType('text/html; charset=utf-8')
+      .send('<html><body>Hi</body></html>');
+    writeResultToResponse(response, response);
+    const result = await observedResponse;
+    expect(result.statusCode).to.equal(200);
+    expect(result.headers).to.have.property(
+      'content-type',
+      'text/html; charset=utf-8',
+    );
+    expect(result.payload).to.equal('<html><body>Hi</body></html>');
+  });
+
+  it('skips writing when the response headers are sent', async () => {
+    response
+      .status(200)
+      .contentType('text/html; charset=utf-8')
+      .send('<html><body>Hi</body></html>');
+    writeResultToResponse(response, undefined);
+    const result = await observedResponse;
+    expect(result.statusCode).to.equal(200);
+    expect(result.headers).to.have.property(
+      'content-type',
+      'text/html; charset=utf-8',
+    );
+    expect(result.payload).to.equal('<html><body>Hi</body></html>');
+  });
+
   function setupResponseMock() {
     const responseMock = stubExpressContext();
     response = responseMock.response;