From a1c701301fdec70db5f0d0d69ceb1a9adb2855ac Mon Sep 17 00:00:00 2001
From: Warren <5959690+wrn14897@users.noreply.github.com>
Date: Sun, 28 Jan 2024 01:45:14 -0800
Subject: [PATCH 1/3] feat: encapsulate shutdown process + use
 http-graceful-shutdown

---
 docker-compose.dev.yml                        |  1 +
 packages/api/package.json                     |  1 +
 .../clickhouse/__tests__/clickhouse.test.ts   |  3 +-
 .../clickhouseGetMultiSeriesChart.test.ts     |  2 +-
 packages/api/src/fixtures.ts                  | 15 ++--
 packages/api/src/index.ts                     | 20 +----
 .../routers/aggregator/__tests__/root.test.ts |  4 +-
 .../src/routers/api/__tests__/alerts.test.ts  |  4 +-
 .../src/routers/api/__tests__/chart.test.ts   |  4 +-
 .../routers/api/__tests__/dashboard.test.ts   |  4 +-
 .../src/routers/api/__tests__/metrics.test.ts |  4 +-
 .../src/routers/api/__tests__/team.test.ts    | 10 +--
 .../external-api/__tests__/alerts.test.ts     |  4 +-
 .../external-api/__tests__/charts.test.ts     |  4 +-
 .../external-api/__tests__/dashboard.test.ts  |  4 +-
 .../routers/external-api/__tests__/v1.test.ts | 10 +--
 packages/api/src/server.ts                    | 84 ++++++++++++-------
 .../src/tasks/__tests__/checkAlerts.test.ts   |  4 +-
 yarn.lock                                     |  7 ++
 19 files changed, 90 insertions(+), 99 deletions(-)

diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml
index 734bb3f62..05b306566 100644
--- a/docker-compose.dev.yml
+++ b/docker-compose.dev.yml
@@ -205,6 +205,7 @@ services:
       CLICKHOUSE_LOG_LEVEL: ${HYPERDX_LOG_LEVEL}
       CLICKHOUSE_PASSWORD: api
       CLICKHOUSE_USER: api
+      DEBUG: 'http-graceful-shutdown'
       EXPRESS_SESSION_SECRET: 'hyperdx is cool 👋'
       FRONTEND_URL: 'http://localhost:8080' # need to be localhost (CORS)
       HDX_NODE_ADVANCED_NETWORK_CAPTURE: 1
diff --git a/packages/api/package.json b/packages/api/package.json
index 083181124..8a16e9ee1 100644
--- a/packages/api/package.json
+++ b/packages/api/package.json
@@ -25,6 +25,7 @@
     "express-session": "^1.17.3",
     "express-winston": "^4.2.0",
     "extract-domain": "^2.4.1",
+    "http-graceful-shutdown": "^3.1.13",
     "isemail": "^3.2.0",
     "jsonwebtoken": "^9.0.0",
     "lodash": "^4.17.21",
diff --git a/packages/api/src/clickhouse/__tests__/clickhouse.test.ts b/packages/api/src/clickhouse/__tests__/clickhouse.test.ts
index 2832d1431..440354c01 100644
--- a/packages/api/src/clickhouse/__tests__/clickhouse.test.ts
+++ b/packages/api/src/clickhouse/__tests__/clickhouse.test.ts
@@ -22,8 +22,7 @@ describe('clickhouse', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   afterEach(async () => {
diff --git a/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts b/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts
index f93f99126..a24613dfe 100644
--- a/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts
+++ b/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts
@@ -20,7 +20,7 @@ describe('clickhouse - getMultiSeriesChart', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
+    await server.stop();
   });
 
   afterEach(async () => {
diff --git a/packages/api/src/fixtures.ts b/packages/api/src/fixtures.ts
index d7c6f36d0..452ed5bf7 100644
--- a/packages/api/src/fixtures.ts
+++ b/packages/api/src/fixtures.ts
@@ -63,6 +63,8 @@ export const initCiEnvs = async () => {
 };
 
 class MockServer extends Server {
+  protected shouldHanleGracefulShutdown = false;
+
   getHttpServer() {
     return this.httpServer;
   }
@@ -75,14 +77,17 @@ class MockServer extends Server {
     await initCiEnvs();
   }
 
-  closeHttpServer() {
+  stop() {
     return new Promise<void>((resolve, reject) => {
       this.httpServer.close(err => {
         if (err) {
           reject(err);
           return;
         }
-        resolve();
+        super
+          .shutdown()
+          .then(() => resolve())
+          .catch(err => reject(err));
       });
     });
   }
@@ -113,10 +118,6 @@ export const getAgent = (server: MockServer) =>
 export const getLoggedInAgent = async (server: MockServer) => {
   const agent = getAgent(server);
 
-  await agent
-    .post('/register/password')
-    .send({ ...MOCK_USER, confirmPassword: 'wrong-password' })
-    .expect(400);
   await agent
     .post('/register/password')
     .send({ ...MOCK_USER, confirmPassword: MOCK_USER.password })
@@ -129,8 +130,6 @@ export const getLoggedInAgent = async (server: MockServer) => {
     throw Error('team or user not found');
   }
 
-  await user.save();
-
   // login app
   await agent.post('/login/password').send(MOCK_USER).expect(302);
 
diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts
index 10154d28d..6070edaf2 100644
--- a/packages/api/src/index.ts
+++ b/packages/api/src/index.ts
@@ -1,9 +1,9 @@
 import { serializeError } from 'serialize-error';
 
-import * as config from './config';
-import Server from './server';
-import { isOperationalError } from './utils/errors';
-import logger from './utils/logger';
+import * as config from '@/config';
+import Server from '@/server';
+import { isOperationalError } from '@/utils/errors';
+import logger from '@/utils/logger';
 
 const server = new Server();
 
@@ -22,16 +22,4 @@ process.on('unhandledRejection', (err: any) => {
   logger.error(serializeError(err));
 });
 
-// graceful shutdown
-process.on('SIGTERM', () => {
-  logger.info('SIGTERM signal received.');
-
-  if (config.IS_DEV) {
-    logger.info('Http server is forced to stop immediately.');
-    process.exit(0);
-  }
-
-  server.stop();
-});
-
 server.start().catch(e => logger.error(serializeError(e)));
diff --git a/packages/api/src/routers/aggregator/__tests__/root.test.ts b/packages/api/src/routers/aggregator/__tests__/root.test.ts
index 21a0a3f19..08d26f928 100644
--- a/packages/api/src/routers/aggregator/__tests__/root.test.ts
+++ b/packages/api/src/routers/aggregator/__tests__/root.test.ts
@@ -5,7 +5,6 @@ import { createTeam } from '@/controllers/team';
 import {
   clearClickhouseTables,
   clearDBCollections,
-  closeDB,
   getAgent,
   getServer,
 } from '@/fixtures';
@@ -24,8 +23,7 @@ describe('aggregator root router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('GET /health', async () => {
diff --git a/packages/api/src/routers/api/__tests__/alerts.test.ts b/packages/api/src/routers/api/__tests__/alerts.test.ts
index e5db993ba..4ee01a702 100644
--- a/packages/api/src/routers/api/__tests__/alerts.test.ts
+++ b/packages/api/src/routers/api/__tests__/alerts.test.ts
@@ -1,6 +1,5 @@
 import {
   clearDBCollections,
-  closeDB,
   getLoggedInAgent,
   getServer,
   makeAlert,
@@ -25,8 +24,7 @@ describe('alerts router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('has alerts attached to dashboards', async () => {
diff --git a/packages/api/src/routers/api/__tests__/chart.test.ts b/packages/api/src/routers/api/__tests__/chart.test.ts
index 0b44f7456..4eb9793c3 100644
--- a/packages/api/src/routers/api/__tests__/chart.test.ts
+++ b/packages/api/src/routers/api/__tests__/chart.test.ts
@@ -5,7 +5,6 @@ import {
   clearClickhouseTables,
   clearDBCollections,
   clearRedis,
-  closeDB,
   generateBuildTeamEventFn,
   getLoggedInAgent,
   getServer,
@@ -26,8 +25,7 @@ describe('charts router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('GET /chart/services', async () => {
diff --git a/packages/api/src/routers/api/__tests__/dashboard.test.ts b/packages/api/src/routers/api/__tests__/dashboard.test.ts
index b7031fa5a..5d1d52a78 100644
--- a/packages/api/src/routers/api/__tests__/dashboard.test.ts
+++ b/packages/api/src/routers/api/__tests__/dashboard.test.ts
@@ -1,6 +1,5 @@
 import {
   clearDBCollections,
-  closeDB,
   getLoggedInAgent,
   getServer,
   makeAlert,
@@ -25,8 +24,7 @@ describe('dashboard router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('deletes attached alerts when deleting charts', async () => {
diff --git a/packages/api/src/routers/api/__tests__/metrics.test.ts b/packages/api/src/routers/api/__tests__/metrics.test.ts
index 23f4df4e7..fb669627a 100644
--- a/packages/api/src/routers/api/__tests__/metrics.test.ts
+++ b/packages/api/src/routers/api/__tests__/metrics.test.ts
@@ -4,7 +4,6 @@ import {
   clearClickhouseTables,
   clearDBCollections,
   clearRedis,
-  closeDB,
   getLoggedInAgent,
   getServer,
 } from '@/fixtures';
@@ -23,8 +22,7 @@ describe('metrics router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('GET /metrics/tags', async () => {
diff --git a/packages/api/src/routers/api/__tests__/team.test.ts b/packages/api/src/routers/api/__tests__/team.test.ts
index 281ac4d1f..fc58eb7a4 100644
--- a/packages/api/src/routers/api/__tests__/team.test.ts
+++ b/packages/api/src/routers/api/__tests__/team.test.ts
@@ -1,11 +1,6 @@
 import _ from 'lodash';
 
-import {
-  clearDBCollections,
-  closeDB,
-  getLoggedInAgent,
-  getServer,
-} from '@/fixtures';
+import { clearDBCollections, getLoggedInAgent, getServer } from '@/fixtures';
 
 describe('team router', () => {
   const server = getServer();
@@ -19,8 +14,7 @@ describe('team router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('GET /team', async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/alerts.test.ts b/packages/api/src/routers/external-api/__tests__/alerts.test.ts
index bfefa78ef..9105f6d87 100644
--- a/packages/api/src/routers/external-api/__tests__/alerts.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/alerts.test.ts
@@ -2,7 +2,6 @@ import _ from 'lodash';
 
 import {
   clearDBCollections,
-  closeDB,
   getLoggedInAgent,
   getServer,
   makeChart,
@@ -33,8 +32,7 @@ describe('/api/v1/alerts', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('CRUD Dashboard Alerts', async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/charts.test.ts b/packages/api/src/routers/external-api/__tests__/charts.test.ts
index 89fc72a77..3c6645ad4 100644
--- a/packages/api/src/routers/external-api/__tests__/charts.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/charts.test.ts
@@ -4,7 +4,6 @@ import ms from 'ms';
 import * as clickhouse from '@/clickhouse';
 import {
   clearDBCollections,
-  closeDB,
   generateBuildTeamEventFn,
   getLoggedInAgent,
   getServer,
@@ -23,8 +22,7 @@ describe('/api/v1/charts/series', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   const now = new Date('2022-01-05').getTime();
diff --git a/packages/api/src/routers/external-api/__tests__/dashboard.test.ts b/packages/api/src/routers/external-api/__tests__/dashboard.test.ts
index 71277af09..bd9ca8ede 100644
--- a/packages/api/src/routers/external-api/__tests__/dashboard.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/dashboard.test.ts
@@ -2,7 +2,6 @@ import _ from 'lodash';
 
 import {
   clearDBCollections,
-  closeDB,
   getLoggedInAgent,
   getServer,
   makeExternalAlert,
@@ -41,8 +40,7 @@ describe('dashboard router', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('CRUD /dashboards', async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/v1.test.ts b/packages/api/src/routers/external-api/__tests__/v1.test.ts
index 2e34a4c98..52987e33c 100644
--- a/packages/api/src/routers/external-api/__tests__/v1.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/v1.test.ts
@@ -1,10 +1,5 @@
 import * as clickhouse from '@/clickhouse';
-import {
-  clearDBCollections,
-  closeDB,
-  getLoggedInAgent,
-  getServer,
-} from '@/fixtures';
+import { clearDBCollections, getLoggedInAgent, getServer } from '@/fixtures';
 
 describe('external api v1', () => {
   const server = getServer();
@@ -19,8 +14,7 @@ describe('external api v1', () => {
   });
 
   afterAll(async () => {
-    await server.closeHttpServer();
-    await closeDB();
+    await server.stop();
   });
 
   it('GET /api/v1', async () => {
diff --git a/packages/api/src/server.ts b/packages/api/src/server.ts
index e694376d4..1800c4354 100644
--- a/packages/api/src/server.ts
+++ b/packages/api/src/server.ts
@@ -1,4 +1,5 @@
 import http from 'http';
+import gracefulShutdown from 'http-graceful-shutdown';
 import { serializeError } from 'serialize-error';
 
 import * as clickhouse from './clickhouse';
@@ -10,6 +11,8 @@ import redisClient from './utils/redis';
 export default class Server {
   protected readonly appType = config.APP_TYPE;
 
+  protected shouldHandleGracefulShutdown = true;
+
   protected httpServer!: http.Server;
 
   private async createServer() {
@@ -29,6 +32,42 @@ export default class Server {
     }
   }
 
+  protected async shutdown(signal?: string) {
+    let hasError = false;
+    logger.info('Closing all db clients...');
+    const [redisCloseResult, mongoCloseResult, clickhouseCloseResult] =
+      await Promise.allSettled([
+        redisClient.disconnect(),
+        mongooseConnection.close(false),
+        clickhouse.client.close(),
+      ]);
+
+    if (redisCloseResult.status === 'rejected') {
+      hasError = true;
+      logger.error(serializeError(redisCloseResult.reason));
+    } else {
+      logger.info('Redis client closed.');
+    }
+
+    if (mongoCloseResult.status === 'rejected') {
+      hasError = true;
+      logger.error(serializeError(mongoCloseResult.reason));
+    } else {
+      logger.info('MongoDB client closed.');
+    }
+
+    if (clickhouseCloseResult.status === 'rejected') {
+      hasError = true;
+      logger.error(serializeError(clickhouseCloseResult.reason));
+    } else {
+      logger.info('Clickhouse client closed.');
+    }
+
+    if (hasError) {
+      throw new Error('Failed to close all clients.');
+    }
+  }
+
   async start() {
     this.httpServer = await this.createServer();
     this.httpServer.keepAliveTimeout = 61000; // Ensure all inactive connections are terminated by the ALB, by setting this a few seconds higher than the ALB idle timeout
@@ -40,39 +79,26 @@ export default class Server {
       );
     });
 
+    if (this.shouldHandleGracefulShutdown) {
+      gracefulShutdown(this.httpServer, {
+        signals: 'SIGINT SIGTERM',
+        timeout: 10000, // 10 secs
+        development: config.IS_DEV,
+        forceExit: true, // triggers process.exit() at the end of shutdown process
+        preShutdown: async () => {
+          console.log('preShutdown...');
+        }, // needed operation before httpConnections are shutted down
+        onShutdown: this.shutdown,
+        finally: () => {
+          logger.info('Server gracefulls shutted down...');
+        }, // finally function (sync) - e.g. for logging
+      });
+    }
+
     await Promise.all([
       connectDB(),
       redisClient.connect(),
       clickhouse.connect(),
     ]);
   }
-
-  // graceful shutdown
-  stop() {
-    this.httpServer.close(closeServerErr => {
-      if (closeServerErr) {
-        logger.error(serializeError(closeServerErr));
-      }
-      logger.info('Http server closed.');
-      redisClient
-        .disconnect()
-        .then(() => {
-          logger.info('Redis client disconnected.');
-        })
-        .catch((err: any) => {
-          logger.error(serializeError(err));
-        });
-      mongooseConnection.close(false, closeDBConnectionErr => {
-        if (closeDBConnectionErr) {
-          logger.error(serializeError(closeDBConnectionErr));
-        }
-        logger.info('Mongo connection closed.');
-
-        if (closeServerErr || closeDBConnectionErr) {
-          process.exit(1);
-        }
-        process.exit(0);
-      });
-    });
-  }
 }
diff --git a/packages/api/src/tasks/__tests__/checkAlerts.test.ts b/packages/api/src/tasks/__tests__/checkAlerts.test.ts
index 7e28f83ee..6a4ad627c 100644
--- a/packages/api/src/tasks/__tests__/checkAlerts.test.ts
+++ b/packages/api/src/tasks/__tests__/checkAlerts.test.ts
@@ -3,7 +3,6 @@ import ms from 'ms';
 import {
   buildMetricSeries,
   clearDBCollections,
-  closeDB,
   generateBuildTeamEventFn,
   getServer,
   mockLogsPropertyTypeMappingsModel,
@@ -132,8 +131,7 @@ describe('checkAlerts', () => {
     });
 
     afterAll(async () => {
-      await server.closeHttpServer();
-      await closeDB();
+      await server.stop();
     });
 
     it('LOG alert', async () => {
diff --git a/yarn.lock b/yarn.lock
index 6c3ac29a6..8423389fb 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -9421,6 +9421,13 @@ http-errors@2.0.0:
     statuses "2.0.1"
     toidentifier "1.0.1"
 
+http-graceful-shutdown@^3.1.13:
+  version "3.1.13"
+  resolved "https://registry.yarnpkg.com/http-graceful-shutdown/-/http-graceful-shutdown-3.1.13.tgz#cf3c8f99787d1f5ac2a6bf8a2132ff54c9ce031e"
+  integrity sha512-Ci5LRufQ8AtrQ1U26AevS8QoMXDOhnAHCJI3eZu1com7mZGHxREmw3dNj85ftpQokQCvak8nI2pnFS8zyM1M+Q==
+  dependencies:
+    debug "^4.3.4"
+
 http-proxy-agent@^4.0.1:
   version "4.0.1"
   resolved "https://registry.yarnpkg.com/http-proxy-agent/-/http-proxy-agent-4.0.1.tgz#8a8c8ef7f5932ccf953c296ca8291b95aa74aa3a"

From 16afd5366cb86f36e040e3bded2345f5899dc4a6 Mon Sep 17 00:00:00 2001
From: Warren <5959690+wrn14897@users.noreply.github.com>
Date: Mon, 29 Jan 2024 10:30:08 -0800
Subject: [PATCH 2/3] style: cleanup

---
 docker-compose.dev.yml       | 1 -
 packages/api/src/fixtures.ts | 2 +-
 packages/api/src/server.ts   | 4 ++--
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml
index 05b306566..734bb3f62 100644
--- a/docker-compose.dev.yml
+++ b/docker-compose.dev.yml
@@ -205,7 +205,6 @@ services:
       CLICKHOUSE_LOG_LEVEL: ${HYPERDX_LOG_LEVEL}
       CLICKHOUSE_PASSWORD: api
       CLICKHOUSE_USER: api
-      DEBUG: 'http-graceful-shutdown'
       EXPRESS_SESSION_SECRET: 'hyperdx is cool 👋'
       FRONTEND_URL: 'http://localhost:8080' # need to be localhost (CORS)
       HDX_NODE_ADVANCED_NETWORK_CAPTURE: 1
diff --git a/packages/api/src/fixtures.ts b/packages/api/src/fixtures.ts
index 452ed5bf7..26f729026 100644
--- a/packages/api/src/fixtures.ts
+++ b/packages/api/src/fixtures.ts
@@ -63,7 +63,7 @@ export const initCiEnvs = async () => {
 };
 
 class MockServer extends Server {
-  protected shouldHanleGracefulShutdown = false;
+  protected shouldHandleGracefulShutdown = false;
 
   getHttpServer() {
     return this.httpServer;
diff --git a/packages/api/src/server.ts b/packages/api/src/server.ts
index 1800c4354..5ef9535a7 100644
--- a/packages/api/src/server.ts
+++ b/packages/api/src/server.ts
@@ -86,8 +86,8 @@ export default class Server {
         development: config.IS_DEV,
         forceExit: true, // triggers process.exit() at the end of shutdown process
         preShutdown: async () => {
-          console.log('preShutdown...');
-        }, // needed operation before httpConnections are shutted down
+          // needed operation before httpConnections are shutted down
+        },
         onShutdown: this.shutdown,
         finally: () => {
           logger.info('Server gracefulls shutted down...');

From c7313a25073e1564da15b9e40e43f17608ed233f Mon Sep 17 00:00:00 2001
From: Warren <5959690+wrn14897@users.noreply.github.com>
Date: Mon, 29 Jan 2024 10:40:44 -0800
Subject: [PATCH 3/3] test: use server.clearDBs in afterEach

---
 .../api/src/clickhouse/__tests__/clickhouse.test.ts |  4 +---
 .../__tests__/clickhouseGetMultiSeriesChart.test.ts |  8 ++------
 packages/api/src/fixtures.ts                        |  8 ++++++++
 .../src/routers/aggregator/__tests__/root.test.ts   | 10 ++--------
 .../api/src/routers/api/__tests__/alerts.test.ts    | 10 ++--------
 .../api/src/routers/api/__tests__/chart.test.ts     |  7 +------
 .../api/src/routers/api/__tests__/dashboard.test.ts | 10 ++--------
 .../api/src/routers/api/__tests__/metrics.test.ts   | 13 ++-----------
 packages/api/src/routers/api/__tests__/team.test.ts |  4 ++--
 .../routers/external-api/__tests__/alerts.test.ts   |  3 +--
 .../routers/external-api/__tests__/charts.test.ts   |  3 +--
 .../external-api/__tests__/dashboard.test.ts        |  3 +--
 .../src/routers/external-api/__tests__/v1.test.ts   |  4 ++--
 .../api/src/tasks/__tests__/checkAlerts.test.ts     |  3 +--
 14 files changed, 28 insertions(+), 62 deletions(-)

diff --git a/packages/api/src/clickhouse/__tests__/clickhouse.test.ts b/packages/api/src/clickhouse/__tests__/clickhouse.test.ts
index 440354c01..a6729ec37 100644
--- a/packages/api/src/clickhouse/__tests__/clickhouse.test.ts
+++ b/packages/api/src/clickhouse/__tests__/clickhouse.test.ts
@@ -3,8 +3,6 @@ import ms from 'ms';
 
 import {
   buildMetricSeries,
-  clearClickhouseTables,
-  closeDB,
   generateBuildTeamEventFn,
   getServer,
   mockLogsPropertyTypeMappingsModel,
@@ -26,7 +24,7 @@ describe('clickhouse', () => {
   });
 
   afterEach(async () => {
-    await clearClickhouseTables();
+    await server.clearDBs();
     jest.clearAllMocks();
   });
 
diff --git a/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts b/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts
index a24613dfe..831737b41 100644
--- a/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts
+++ b/packages/api/src/clickhouse/__tests__/clickhouseGetMultiSeriesChart.test.ts
@@ -2,11 +2,7 @@ import _ from 'lodash';
 import ms from 'ms';
 
 import * as clickhouse from '@/clickhouse';
-import {
-  buildMetricSeries,
-  clearClickhouseTables,
-  getServer,
-} from '@/fixtures';
+import { buildMetricSeries, getServer } from '@/fixtures';
 
 describe('clickhouse - getMultiSeriesChart', () => {
   const server = getServer();
@@ -24,7 +20,7 @@ describe('clickhouse - getMultiSeriesChart', () => {
   });
 
   afterEach(async () => {
-    await clearClickhouseTables();
+    await server.clearDBs();
     jest.clearAllMocks();
   });
 
diff --git a/packages/api/src/fixtures.ts b/packages/api/src/fixtures.ts
index 26f729026..32bb9ef22 100644
--- a/packages/api/src/fixtures.ts
+++ b/packages/api/src/fixtures.ts
@@ -91,6 +91,14 @@ class MockServer extends Server {
       });
     });
   }
+
+  clearDBs() {
+    return Promise.all([
+      clearDBCollections(),
+      clearClickhouseTables(),
+      clearRedis(),
+    ]);
+  }
 }
 
 class MockAPIServer extends MockServer {
diff --git a/packages/api/src/routers/aggregator/__tests__/root.test.ts b/packages/api/src/routers/aggregator/__tests__/root.test.ts
index 08d26f928..adcc95569 100644
--- a/packages/api/src/routers/aggregator/__tests__/root.test.ts
+++ b/packages/api/src/routers/aggregator/__tests__/root.test.ts
@@ -2,12 +2,7 @@ import _ from 'lodash';
 
 import * as clickhouse from '@/clickhouse';
 import { createTeam } from '@/controllers/team';
-import {
-  clearClickhouseTables,
-  clearDBCollections,
-  getAgent,
-  getServer,
-} from '@/fixtures';
+import { getAgent, getServer } from '@/fixtures';
 import { sleep } from '@/utils/common';
 
 describe('aggregator root router', () => {
@@ -18,8 +13,7 @@ describe('aggregator root router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
-    await clearClickhouseTables();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/api/__tests__/alerts.test.ts b/packages/api/src/routers/api/__tests__/alerts.test.ts
index 4ee01a702..53e863e11 100644
--- a/packages/api/src/routers/api/__tests__/alerts.test.ts
+++ b/packages/api/src/routers/api/__tests__/alerts.test.ts
@@ -1,10 +1,4 @@
-import {
-  clearDBCollections,
-  getLoggedInAgent,
-  getServer,
-  makeAlert,
-  makeChart,
-} from '@/fixtures';
+import { getLoggedInAgent, getServer, makeAlert, makeChart } from '@/fixtures';
 
 const MOCK_DASHBOARD = {
   name: 'Test Dashboard',
@@ -20,7 +14,7 @@ describe('alerts router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/api/__tests__/chart.test.ts b/packages/api/src/routers/api/__tests__/chart.test.ts
index 4eb9793c3..269b91aa5 100644
--- a/packages/api/src/routers/api/__tests__/chart.test.ts
+++ b/packages/api/src/routers/api/__tests__/chart.test.ts
@@ -2,9 +2,6 @@ import ms from 'ms';
 
 import * as clickhouse from '@/clickhouse';
 import {
-  clearClickhouseTables,
-  clearDBCollections,
-  clearRedis,
   generateBuildTeamEventFn,
   getLoggedInAgent,
   getServer,
@@ -19,9 +16,7 @@ describe('charts router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
-    await clearClickhouseTables();
-    await clearRedis();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/api/__tests__/dashboard.test.ts b/packages/api/src/routers/api/__tests__/dashboard.test.ts
index 5d1d52a78..bcca681a4 100644
--- a/packages/api/src/routers/api/__tests__/dashboard.test.ts
+++ b/packages/api/src/routers/api/__tests__/dashboard.test.ts
@@ -1,10 +1,4 @@
-import {
-  clearDBCollections,
-  getLoggedInAgent,
-  getServer,
-  makeAlert,
-  makeChart,
-} from '@/fixtures';
+import { getLoggedInAgent, getServer, makeAlert, makeChart } from '@/fixtures';
 
 const MOCK_DASHBOARD = {
   name: 'Test Dashboard',
@@ -20,7 +14,7 @@ describe('dashboard router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/api/__tests__/metrics.test.ts b/packages/api/src/routers/api/__tests__/metrics.test.ts
index fb669627a..8019c1249 100644
--- a/packages/api/src/routers/api/__tests__/metrics.test.ts
+++ b/packages/api/src/routers/api/__tests__/metrics.test.ts
@@ -1,12 +1,5 @@
 import * as clickhouse from '@/clickhouse';
-import {
-  buildMetricSeries,
-  clearClickhouseTables,
-  clearDBCollections,
-  clearRedis,
-  getLoggedInAgent,
-  getServer,
-} from '@/fixtures';
+import { buildMetricSeries, getLoggedInAgent, getServer } from '@/fixtures';
 
 describe('metrics router', () => {
   const server = getServer();
@@ -16,9 +9,7 @@ describe('metrics router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
-    await clearClickhouseTables();
-    await clearRedis();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/api/__tests__/team.test.ts b/packages/api/src/routers/api/__tests__/team.test.ts
index fc58eb7a4..057052906 100644
--- a/packages/api/src/routers/api/__tests__/team.test.ts
+++ b/packages/api/src/routers/api/__tests__/team.test.ts
@@ -1,6 +1,6 @@
 import _ from 'lodash';
 
-import { clearDBCollections, getLoggedInAgent, getServer } from '@/fixtures';
+import { getLoggedInAgent, getServer } from '@/fixtures';
 
 describe('team router', () => {
   const server = getServer();
@@ -10,7 +10,7 @@ describe('team router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/alerts.test.ts b/packages/api/src/routers/external-api/__tests__/alerts.test.ts
index 9105f6d87..9780ed59f 100644
--- a/packages/api/src/routers/external-api/__tests__/alerts.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/alerts.test.ts
@@ -1,7 +1,6 @@
 import _ from 'lodash';
 
 import {
-  clearDBCollections,
   getLoggedInAgent,
   getServer,
   makeChart,
@@ -28,7 +27,7 @@ describe('/api/v1/alerts', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/charts.test.ts b/packages/api/src/routers/external-api/__tests__/charts.test.ts
index 3c6645ad4..2dcbfa36d 100644
--- a/packages/api/src/routers/external-api/__tests__/charts.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/charts.test.ts
@@ -3,7 +3,6 @@ import ms from 'ms';
 
 import * as clickhouse from '@/clickhouse';
 import {
-  clearDBCollections,
   generateBuildTeamEventFn,
   getLoggedInAgent,
   getServer,
@@ -18,7 +17,7 @@ describe('/api/v1/charts/series', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/dashboard.test.ts b/packages/api/src/routers/external-api/__tests__/dashboard.test.ts
index bd9ca8ede..f32a5e7b3 100644
--- a/packages/api/src/routers/external-api/__tests__/dashboard.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/dashboard.test.ts
@@ -1,7 +1,6 @@
 import _ from 'lodash';
 
 import {
-  clearDBCollections,
   getLoggedInAgent,
   getServer,
   makeExternalAlert,
@@ -36,7 +35,7 @@ describe('dashboard router', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
   });
 
   afterAll(async () => {
diff --git a/packages/api/src/routers/external-api/__tests__/v1.test.ts b/packages/api/src/routers/external-api/__tests__/v1.test.ts
index 52987e33c..644f9379e 100644
--- a/packages/api/src/routers/external-api/__tests__/v1.test.ts
+++ b/packages/api/src/routers/external-api/__tests__/v1.test.ts
@@ -1,5 +1,5 @@
 import * as clickhouse from '@/clickhouse';
-import { clearDBCollections, getLoggedInAgent, getServer } from '@/fixtures';
+import { getLoggedInAgent, getServer } from '@/fixtures';
 
 describe('external api v1', () => {
   const server = getServer();
@@ -9,7 +9,7 @@ describe('external api v1', () => {
   });
 
   afterEach(async () => {
-    await clearDBCollections();
+    await server.clearDBs();
     jest.clearAllMocks();
   });
 
diff --git a/packages/api/src/tasks/__tests__/checkAlerts.test.ts b/packages/api/src/tasks/__tests__/checkAlerts.test.ts
index 6a4ad627c..fe275c724 100644
--- a/packages/api/src/tasks/__tests__/checkAlerts.test.ts
+++ b/packages/api/src/tasks/__tests__/checkAlerts.test.ts
@@ -2,7 +2,6 @@ import ms from 'ms';
 
 import {
   buildMetricSeries,
-  clearDBCollections,
   generateBuildTeamEventFn,
   getServer,
   mockLogsPropertyTypeMappingsModel,
@@ -126,7 +125,7 @@ describe('checkAlerts', () => {
     });
 
     afterEach(async () => {
-      await clearDBCollections();
+      await server.clearDBs();
       jest.clearAllMocks();
     });