CM-11314 Add integrational test for getTracesCntForWorkspaces endpoint

comet-ml · Sep 11, 2024 · 9c460dc · 9c460dc
1 parent 3948b12
commit 9c460dc
Showing 1 changed file with 47 additions and 13 deletions.
diff --git a/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java b/apps/opik-backend/src/test/java/com/comet/opik/api/resources/v1/priv/TracesResourceTest.java
@@ -7,6 +7,8 @@
 import com.comet.opik.api.Project;
 import com.comet.opik.api.ScoreSource;
 import com.comet.opik.api.Trace;
+import com.comet.opik.api.TraceCntRequest;
+import com.comet.opik.api.TraceCntResponse;
 import com.comet.opik.api.TraceUpdate;
 import com.comet.opik.api.error.ErrorMessage;
 import com.comet.opik.api.filter.Filter;
@@ -315,18 +317,7 @@ void get__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean e
             var workspaceName = UUID.randomUUID().toString();
             var workspaceId = UUID.randomUUID().toString();
 
-            mockTargetWorkspace(okApikey, workspaceName, workspaceId);
-
-            var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class)
-                    .stream()
-                    .map(t -> t.toBuilder()
-                            .projectId(null)
-                            .projectName(DEFAULT_PROJECT)
-                            .feedbackScores(null)
-                            .build())
-                    .toList();
-
-            traces.forEach(trace -> TracesResourceTest.this.create(trace, okApikey, workspaceName));
+            int tracesCnt = setupTracesForWorkspace(workspaceName, workspaceId);
 
             try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI))
                     .queryParam("project_name", DEFAULT_PROJECT)
@@ -340,7 +331,7 @@ void get__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean e
                     assertThat(actualResponse.hasEntity()).isTrue();
 
                     var response = actualResponse.readEntity(Trace.TracePage.class);
-                    assertThat(response.content()).hasSize(traces.size());
+                    assertThat(response.content()).hasSize(tracesCnt);
                 } else {
                     assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(401);
                     assertThat(actualResponse.readEntity(io.dropwizard.jersey.errors.ErrorMessage.class))
@@ -349,6 +340,32 @@ void get__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey, boolean e
             }
         }
 
+        @Test
+        @DisplayName("Get traces count for each workspace from requested list, no Auth")
+        void tracesCountForWorkspace() {
+            var workspaceName = UUID.randomUUID().toString();
+            var workspaceId = UUID.randomUUID().toString();
+
+            int tracesCnt = setupTracesForWorkspace(workspaceName, workspaceId);
+            TraceCntRequest request = new TraceCntRequest(List.of(workspaceId));
+
+            try (var actualResponse = client.target(URL_TEMPLATE.formatted(baseURI))
+                    .path("/workspaces-traces")
+                    .request()
+                    .header(HttpHeaders.AUTHORIZATION, okApikey)
+                    .header(WORKSPACE_HEADER, workspaceName)
+                    .post(Entity.json(request))) {
+
+                assertThat(actualResponse.getStatusInfo().getStatusCode()).isEqualTo(200);
+                assertThat(actualResponse.hasEntity()).isTrue();
+
+                var response = actualResponse.readEntity(TraceCntResponse.class);
+                assertThat(response.workspacesTracesCnt()).hasSize(1);
+                assertThat(response.workspacesTracesCnt().get(0).workspace()).isEqualTo(workspaceId);
+                assertThat(response.workspacesTracesCnt().get(0).tracesCnt()).isEqualTo(tracesCnt);
+            }
+        }
+
         @ParameterizedTest
         @MethodSource("credentials")
         @DisplayName("Trace feedback, when api key is present, then return proper response")
@@ -460,6 +477,23 @@ void feedbackBatch__whenApiKeyIsPresent__thenReturnProperResponse(String apiKey,
 
         }
 
+        private int setupTracesForWorkspace(String workspaceName, String workspaceId) {
+            mockTargetWorkspace(okApikey, workspaceName, workspaceId);
+
+            var traces = PodamFactoryUtils.manufacturePojoList(factory, Trace.class)
+                    .stream()
+                    .map(t -> t.toBuilder()
+                            .projectId(null)
+                            .projectName(DEFAULT_PROJECT)
+                            .feedbackScores(null)
+                            .build())
+                    .toList();
+
+            traces.forEach(trace -> TracesResourceTest.this.create(trace, okApikey, workspaceName));
+
+            return traces.size();
+        }
+
     }
 
     @Nested