test

Author: liustve

validator/src/main/java/com/amazon/aoc/validators/TraceValidator.java

@@ -149,23 +149,13 @@ private Map<String, Object> getTrace() throws Exception {
     }
 
     if (validationConfig.getHttpPath().contains("chat")) {
-      // For chat endpoints, directly use the trace ID instead of searching
-      List<String> traceIdLists = Collections.singletonList(context.getTraceId());
-      List<Trace> retrievedTraceList = xrayService.listTraceByIds(traceIdLists);
-      if (retrievedTraceList == null || retrievedTraceList.isEmpty()) {
-        throw new BaseException(ExceptionCode.EMPTY_LIST);
-      }
-      return this.flattenDocument(retrievedTraceList.get(0).getSegments());
+      return this.getTraceById(Collections.singletonList(context.getTraceId()));
     }
+    
     log.info("Trace Filter: {}", traceFilter);
     List<TraceSummary> retrieveTraceLists = xrayService.searchTraces(traceFilter);
     List<String> traceIdLists = Collections.singletonList(retrieveTraceLists.get(0).getId());
-    List<Trace> retrievedTraceList = xrayService.listTraceByIds(traceIdLists);
-
-    if (retrievedTraceList == null || retrievedTraceList.isEmpty()) {
-      throw new BaseException(ExceptionCode.EMPTY_LIST);
-    }
-    return this.flattenDocument(retrievedTraceList.get(0).getSegments());
+    return getTraceById(traceIdLists);
   }
 
   private Map<String, Object> flattenDocument(List<Segment> segmentList) {
@@ -200,6 +190,14 @@ private Map<String, Object> flattenDocument(List<Segment> segmentList) {
     return JsonFlattener.flattenAsMap(segmentsJson.toString());
   }
 
+  private Map<String, Object> getTraceById(List<String> traceIdLists) throws Exception {
+    List<Trace> retrievedTraceList = xrayService.listTraceByIds(traceIdLists);
+    if (retrievedTraceList == null || retrievedTraceList.isEmpty()) {
+      throw new BaseException(ExceptionCode.EMPTY_LIST);
+    }
+    return this.flattenDocument(retrievedTraceList.get(0).getSegments());
+  }
+
   // This method will get the stored traces
   private Map<String, Object> getStoredTrace() throws Exception {
     Map<String, Object> flattenedJsonMapForStoredTraces = null;