Fixing the Zipkin API (endpoint: /api/v2/trace/{traceId}) (#1140)

* Fixing the Zipkin API (endpoint: /api/v2/trace/{traceId})

* Adding tests cases

* Formatting fix

* Trigger CI

astra/pom.xml

@@ -249,6 +249,11 @@
         </dependency>
 
         <!-- Misc -->
+        <dependency>
+            <groupId>org.json</groupId>
+            <artifactId>json</artifactId>
+            <version>20210307</version>
+        </dependency>
         <dependency>
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-text</artifactId>

astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java

@@ -33,6 +33,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
+import org.json.JSONObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -200,7 +201,13 @@ public HttpResponse getTraceByTraceId(
       @Param("endTimeEpochMs") Optional<Long> endTimeEpochMs,
       @Param("maxSpans") Optional<Integer> maxSpans)
       throws IOException {
-    String queryString = "trace_id:" + traceId;
+
+    JSONObject traceObject = new JSONObject();
+    traceObject.put("trace_id", traceId);
+    JSONObject queryJson = new JSONObject();
+    queryJson.put("term", traceObject);
+    String queryString = queryJson.toString();
+
     long startTime =
         startTimeEpochMs.orElseGet(
             () -> Instant.now().minus(LOOKBACK_MINS, ChronoUnit.MINUTES).toEpochMilli());

astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java

@@ -0,0 +1,71 @@
+package com.slack.astra.zipkinApi;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import brave.Span;
+import brave.Tracing;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.io.Resources;
+import com.google.protobuf.util.JsonFormat;
+import com.linecorp.armeria.common.AggregatedHttpResponse;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.HttpStatus;
+import com.slack.astra.proto.service.AstraSearch;
+import com.slack.astra.server.AstraQueryServiceBase;
+import java.io.IOException;
+import java.util.Optional;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+import org.mockito.MockedStatic;
+import org.mockito.MockitoAnnotations;
+
+public class ZipkinServiceTest {
+  @Mock private AstraQueryServiceBase searcher;
+  private ZipkinService zipkinService;
+  private AstraSearch.SearchResult mockSearchResult;
+
+  @BeforeEach
+  public void setup() throws IOException {
+    MockitoAnnotations.openMocks(this);
+    zipkinService = spy(new ZipkinService(searcher));
+    // Build mockSearchResult
+    ObjectMapper objectMapper = new ObjectMapper();
+    JsonNode jsonNode =
+        objectMapper.readTree(Resources.getResource("zipkinApi/search_result.json"));
+    String jsonString = objectMapper.writeValueAsString(jsonNode);
+    AstraSearch.SearchResult.Builder builder = AstraSearch.SearchResult.newBuilder();
+    JsonFormat.parser().merge(jsonString, builder);
+    mockSearchResult = builder.build();
+  }
+
+  @Test
+  public void testGetTraceByTraceId_onlyTraceIdProvided() throws Exception {
+
+    try (MockedStatic<Tracing> mockedTracing = mockStatic(Tracing.class)) {
+      brave.Tracer mockTracer = mock(brave.Tracer.class);
+      Span mockSpan = mock(Span.class);
+
+      mockedTracing.when(Tracing::currentTracer).thenReturn(mockTracer);
+      when(mockTracer.currentSpan()).thenReturn(mockSpan);
+      String traceId = "test_trace_1";
+
+      when(searcher.doSearch(any())).thenReturn(mockSearchResult);
+
+      // Act
+      HttpResponse response =
+          zipkinService.getTraceByTraceId(
+              traceId, Optional.empty(), Optional.empty(), Optional.empty());
+      AggregatedHttpResponse aggregatedResponse = response.aggregate().join();
+
+      // Assert
+      assertEquals(HttpStatus.OK, aggregatedResponse.status());
+    }
+  }
+}

astra/src/test/resources/zipkinApi/search_result.json

@@ -0,0 +1,10 @@
+{
+  "hits": [
+    "{\"id\":\"101\",\"timestamp\":1730407713.560000000,\"source\":{\"service-name\":\"traces-staging\",\"trace_id\":\"1234556789\",\"_timesinceepoch\":\"2024-10-31T20:48:33.560Z\",\"level\":\"INFO\",\"service_name\":\"all_tenants\",\"message\":\"This is a test log message\",\"timestamp\":\"2024-03-20T15:30:00.000Z\"},\"index\":\"all_tenants\"}",
+    "{\"id\":\"100\",\"timestamp\":1730407435.071000000,\"source\":{\"service-name\":\"traces-staging\",\"trace_id\":\"1234556789\",\"_timesinceepoch\":\"2024-10-31T20:43:55.071Z\",\"level\":\"INFO\",\"service_name\":\"all_tenants\",\"message\":\"This is a test log message 1\",\"timestamp\":\"2024-03-20T15:30:00.000Z\"},\"index\":\"all_tenants\"}"
+  ],
+  "tookMicros": "2732",
+  "totalNodes": 1,
+  "totalSnapshots": 1,
+  "snapshotsWithReplicas": 1
+}