From 1a47a4222568703ca16747f75c89ec11b59425dd Mon Sep 17 00:00:00 2001 From: Zarna Parekh Date: Thu, 7 Nov 2024 09:17:37 -0800 Subject: [PATCH] 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 | 5 ++ .../slack/astra/zipkinApi/ZipkinService.java | 9 ++- .../astra/zipkinApi/ZipkinServiceTest.java | 71 +++++++++++++++++++ .../resources/zipkinApi/search_result.json | 10 +++ 4 files changed, 94 insertions(+), 1 deletion(-) create mode 100644 astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java create mode 100644 astra/src/test/resources/zipkinApi/search_result.json diff --git a/astra/pom.xml b/astra/pom.xml index 623b49bd7b..472b52b9f4 100644 --- a/astra/pom.xml +++ b/astra/pom.xml @@ -249,6 +249,11 @@ + + org.json + json + 20210307 + org.apache.commons commons-text diff --git a/astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java b/astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java index c7de18ae1f..4d8420e6a0 100644 --- a/astra/src/main/java/com/slack/astra/zipkinApi/ZipkinService.java +++ b/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 endTimeEpochMs, @Param("maxSpans") Optional 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()); diff --git a/astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java b/astra/src/test/java/com/slack/astra/zipkinApi/ZipkinServiceTest.java new file mode 100644 index 0000000000..9154d28a7f --- /dev/null +++ b/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 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()); + } + } +} diff --git a/astra/src/test/resources/zipkinApi/search_result.json b/astra/src/test/resources/zipkinApi/search_result.json new file mode 100644 index 0000000000..4973af787d --- /dev/null +++ b/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 +} \ No newline at end of file