From 742da55adf4e013275ca535a060a907fc7ea64f9 Mon Sep 17 00:00:00 2001
From: Joshua Li <joshuali925@gmail.com>
Date: Fri, 15 Nov 2024 19:14:00 +0000
Subject: [PATCH] update unit tests

Signed-off-by: Joshua Li <joshuali925@gmail.com>
---
 .../server/routes/index.test.ts               | 21 +++++++++++++++++++
 .../query_enhancements/server/routes/index.ts | 12 +++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 src/plugins/query_enhancements/server/routes/index.test.ts

diff --git a/src/plugins/query_enhancements/server/routes/index.test.ts b/src/plugins/query_enhancements/server/routes/index.test.ts
new file mode 100644
index 00000000000..9c7c7a56de2
--- /dev/null
+++ b/src/plugins/query_enhancements/server/routes/index.test.ts
@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { coerceStatusCode } from '.';
+
+describe('coerceStatusCode', () => {
+  it('should return 503 when input is 500', () => {
+    expect(coerceStatusCode(500)).toBe(503);
+  });
+
+  it('should return the input status code when it is not 500', () => {
+    expect(coerceStatusCode(404)).toBe(404);
+  });
+
+  it('should return 503 when input is undefined or null', () => {
+    expect(coerceStatusCode((undefined as unknown) as number)).toBe(503);
+    expect(coerceStatusCode((null as unknown) as number)).toBe(503);
+  });
+});
diff --git a/src/plugins/query_enhancements/server/routes/index.ts b/src/plugins/query_enhancements/server/routes/index.ts
index eeb8e935c38..84cf19bec50 100644
--- a/src/plugins/query_enhancements/server/routes/index.ts
+++ b/src/plugins/query_enhancements/server/routes/index.ts
@@ -16,6 +16,15 @@ import { API } from '../../common';
 import { registerQueryAssistRoutes } from './query_assist';
 import { registerDataSourceConnectionsRoutes } from './data_source_connection';
 
+/**
+ * Coerce status code to 503 for 500 errors from dependency services. Only use
+ * this function to handle errors throw by other services, and not from OSD.
+ */
+export const coerceStatusCode = (statusCode: number) => {
+  if (statusCode === 500) return 503;
+  return statusCode || 503;
+};
+
 /**
  * @experimental
  *
@@ -91,9 +100,8 @@ export function defineSearchStrategyRouteProvider(logger: Logger, router: IRoute
           } catch (e) {
             error = err;
           }
-          const statusCode = error.status || err.status;
           return res.custom({
-            statusCode: statusCode === 500 ? 503 : statusCode || 503,
+            statusCode: coerceStatusCode(error.status || err.status),
             body: err.message,
           });
         }