fix: Base assistant is shown in the listing API response #103

epam · Dec 28, 2023 · b49b0c3 · b49b0c3
1 parent f8274a3
commit b49b0c3
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 4 deletions.
diff --git a/src/main/java/com/epam/aidial/core/controller/AssistantController.java b/src/main/java/com/epam/aidial/core/controller/AssistantController.java
@@ -12,6 +12,8 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static com.epam.aidial.core.config.Config.ASSISTANT;
+
 @RequiredArgsConstructor
 public class AssistantController {
 
@@ -21,7 +23,7 @@ public Future<?> getAssistant(String assistantId) {
         Config config = context.getConfig();
         Assistant assistant = config.getAssistant().getAssistants().get(assistantId);
 
-        if (assistant == null) {
+        if (assistant == null || ASSISTANT.equals(assistant.getName())) {
             return context.respond(HttpStatus.NOT_FOUND);
         }
 
@@ -38,7 +40,7 @@ public Future<?> getAssistants() {
         List<AssistantData> assistants = new ArrayList<>();
 
         for (Assistant assistant : config.getAssistant().getAssistants().values()) {
-            if (DeploymentController.hasAccess(context, assistant)) {
+            if (!ASSISTANT.equals(assistant.getName()) && DeploymentController.hasAccess(context, assistant)) {
                 AssistantData data = createAssistant(assistant);
                 assistants.add(data);
             }

diff --git a/src/test/java/com/epam/aidial/core/controller/AssistantControllerTest.java b/src/test/java/com/epam/aidial/core/controller/AssistantControllerTest.java
@@ -0,0 +1,99 @@
+package com.epam.aidial.core.controller;
+
+import com.epam.aidial.core.ProxyContext;
+import com.epam.aidial.core.config.Assistant;
+import com.epam.aidial.core.config.Assistants;
+import com.epam.aidial.core.config.Config;
+import com.epam.aidial.core.data.AssistantData;
+import com.epam.aidial.core.data.ListData;
+import com.epam.aidial.core.util.HttpStatus;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentMatcher;
+import org.mockito.InjectMocks;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.util.List;
+
+import static com.epam.aidial.core.config.Config.ASSISTANT;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+public class AssistantControllerTest {
+
+    @Mock
+    private ProxyContext context;
+
+    private Config config;
+
+    @InjectMocks
+    private AssistantController controller;
+
+    @BeforeEach
+    public void beforeEach() {
+        config = new Config();
+        Assistants assistants = new Assistants();
+        Assistant baseAssistant = new Assistant();
+        baseAssistant.setName(ASSISTANT);
+        assistants.getAssistants().put(ASSISTANT, baseAssistant);
+        config.setAssistant(assistants);
+        when(context.getConfig()).thenReturn(config);
+    }
+
+    @Test
+    public void testGetAssistant_Success() {
+        Assistant assistant = new Assistant();
+        assistant.setName("key");
+        config.getAssistant().getAssistants().put("key", assistant);
+
+        controller.getAssistant("key");
+
+        verify(context).respond(eq(HttpStatus.OK), argThat((ArgumentMatcher<Object>) argument -> {
+            if (!(argument instanceof AssistantData data)) {
+                return false;
+            }
+            return "key".equals(data.getId());
+        }));
+    }
+
+    @Test
+    public void testGetAssistant_NotFound() {
+        controller.getAssistant("key");
+
+        verify(context).respond(eq(HttpStatus.NOT_FOUND));
+    }
+
+    @Test
+    public void testGetAssistant_NotFoundBaseAssistant() {
+
+        controller.getAssistant(ASSISTANT);
+
+        verify(context).respond(eq(HttpStatus.NOT_FOUND));
+    }
+
+    @Test
+    public void testGetAssistants_Success() {
+        Assistant assistant = new Assistant();
+        assistant.setName("key");
+        config.getAssistant().getAssistants().put("key", assistant);
+
+        controller.getAssistants();
+
+        verify(context).respond(eq(HttpStatus.OK), argThat((ArgumentMatcher<Object>) argument -> {
+            if (!(argument instanceof ListData list)) {
+                return false;
+            }
+            List<AssistantData> data = list.getData();
+            if (data.size() != 1) {
+                return false;
+            }
+            return "key".equals(data.get(0).getId());
+        }));
+    }
+
+}
diff --git a/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java b/src/test/java/com/epam/aidial/core/controller/DeploymentPostControllerTest.java
@@ -40,12 +40,10 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.assertArg;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;