From 49bba29c7514fd4dc831b5e43691dcaf26afdc17 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Mon, 29 Jan 2024 12:46:01 +0300 Subject: [PATCH 1/2] fix: NPE occurs on rate limit hit in PostDeploymentController #178 --- .../epam/aidial/core/controller/DeploymentPostController.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java b/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java index 065356b76..ecbef54aa 100644 --- a/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java +++ b/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java @@ -540,6 +540,8 @@ private Future respond(HttpStatus status, Object result) { private void finalizeRequest() { proxy.getTokenStatsTracker().endSpan(context); - proxy.getApiKeyStore().invalidateApiKey(context.getProxyApiKeyData()); + if (context.getProxyApiKeyData() != null) { + proxy.getApiKeyStore().invalidateApiKey(context.getProxyApiKeyData()); + } } } \ No newline at end of file From 8b54f5b3258eb14a3a1b857c88522cb1c6073758 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Mon, 29 Jan 2024 13:32:40 +0300 Subject: [PATCH 2/2] fix: fix Unit tests --- .../core/controller/DeploymentPostControllerTest.java | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java b/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java index e690708a0..3bd69be28 100644 --- a/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java +++ b/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java @@ -88,7 +88,6 @@ public class DeploymentPostControllerTest { public void testUnsupportedContentType() { when(context.getRequest()).thenReturn(request); when(request.getHeader(eq(HttpHeaders.CONTENT_TYPE))).thenReturn("unsupported"); - when(proxy.getApiKeyStore()).thenReturn(apiKeyStore); when(proxy.getTokenStatsTracker()).thenReturn(tokenStatsTracker); controller.handle("app1", "api"); @@ -147,7 +146,6 @@ public void testNoRoute() { MultiMap headers = mock(MultiMap.class); when(request.headers()).thenReturn(headers); when(context.getDeployment()).thenReturn(application); - when(proxy.getApiKeyStore()).thenReturn(apiKeyStore); when(proxy.getTokenStatsTracker()).thenReturn(tokenStatsTracker); controller.handle("app1", "chat/completions"); @@ -333,14 +331,12 @@ public void testHandleResponse_Model() { when(context.getUpstreamRoute()).thenReturn(upstreamRoute); when(context.getResponseBody()).thenReturn(Buffer.buffer()); when(proxy.getTokenStatsTracker()).thenReturn(tokenStatsTracker); - when(proxy.getApiKeyStore()).thenReturn(apiKeyStore); controller.handleResponse(); verify(rateLimiter).increase(eq(context)); verify(context).setTokenUsage(any(TokenUsage.class)); verify(logStore).save(eq(context)); - verify(apiKeyStore).invalidateApiKey(any()); verify(tokenStatsTracker).endSpan(eq(context)); } @@ -358,7 +354,6 @@ public void testHandleResponse_App() { when(response.getStatusCode()).thenReturn(HttpStatus.OK.getCode()); when(context.getResponseBody()).thenReturn(Buffer.buffer()); when(proxy.getTokenStatsTracker()).thenReturn(tokenStatsTracker); - when(proxy.getApiKeyStore()).thenReturn(apiKeyStore); when(tokenStatsTracker.getTokenStats(eq(context))).thenReturn(Future.succeededFuture(new TokenUsage())); controller.handleResponse(); @@ -367,7 +362,6 @@ public void testHandleResponse_App() { verify(tokenStatsTracker).getTokenStats(eq(context)); verify(context).setTokenUsage(any(TokenUsage.class)); verify(logStore).save(eq(context)); - verify(apiKeyStore).invalidateApiKey(any()); verify(tokenStatsTracker).endSpan(eq(context)); }