From 38436041d3158a6b5249e848d2a68af68589b315 Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Fri, 13 Dec 2024 10:50:09 +0100
Subject: [PATCH] catalog-backend: add test for response writing and properly
 stop on close

Signed-off-by: Patrik Oldsberg <poldsberg@gmail.com>
---
 .../src/service/response/write.test.ts        | 36 +++++++++++++++++++
 .../src/service/response/write.ts             | 24 ++++++++-----
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/plugins/catalog-backend/src/service/response/write.test.ts b/plugins/catalog-backend/src/service/response/write.test.ts
index b1a688b48d..fdca874904 100644
--- a/plugins/catalog-backend/src/service/response/write.test.ts
+++ b/plugins/catalog-backend/src/service/response/write.test.ts
@@ -277,5 +277,41 @@ describe('writeEntitiesResponse', () => {
         totalItems: 1337,
       });
     });
+
+    it('should write a large wrapped response', async () => {
+      const entityMock = JSON.stringify({
+        apiVersion: 'backstage.io/v1alpha1',
+        kind: 'Component',
+        metadata: {
+          name: 'my-component',
+          namespace: 'default',
+          annotations: {
+            'backstage.io/managed-by-location': 'url:https://example.com',
+          },
+        },
+        spec: {
+          type: 'service',
+          owner: 'me',
+          lifecycle: 'production',
+        },
+      });
+      const res = await request(app)
+        .get('/wrapped')
+        .send({
+          type: 'raw',
+          entities: Array(300).fill(entityMock),
+        });
+
+      expect(res.status).toBe(200);
+      expect(res.type).toBe('application/json');
+      expect(res.header['content-type']).toBe(
+        'application/json; charset=utf-8',
+      );
+      expect(res.body).toEqual({
+        page: 1,
+        items: expect.objectContaining({ length: 300 }),
+        totalItems: 1337,
+      });
+    });
   });
 });
diff --git a/plugins/catalog-backend/src/service/response/write.ts b/plugins/catalog-backend/src/service/response/write.ts
index 0962ee15d1..c2bca0953a 100644
--- a/plugins/catalog-backend/src/service/response/write.ts
+++ b/plugins/catalog-backend/src/service/response/write.ts
@@ -80,15 +80,23 @@ export async function writeEntitiesResponse(
 
     const needsDrain = !res.write(prefix + entity, 'utf8');
     if (needsDrain) {
-      await new Promise<void>(resolve => {
-        const cont = () => {
-          res.off('drain', cont);
-          res.off('close', cont);
-          resolve();
-        };
-        res.on('drain', cont);
-        res.on('close', cont);
+      const closed = await new Promise<boolean>(resolve => {
+        function onContinue() {
+          res.off('drain', onContinue);
+          res.off('close', onClose);
+          resolve(false);
+        }
+        function onClose() {
+          res.off('drain', onContinue);
+          res.off('close', onClose);
+          resolve(true);
+        }
+        res.on('drain', onContinue);
+        res.on('close', onClose);
       });
+      if (closed) {
+        return;
+      }
     }
   }
   res.end(`${first ? '[' : ''}]${trailing}`);