Add transport tests, refactor toXContent tests

Signed-off-by: Andy Qin <[email protected]>

Author: q-andy

src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java

@@ -18,6 +18,8 @@
 import com.google.common.collect.ImmutableList;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.neuralsearch.settings.NeuralSearchSettingsAccessor;
+import org.opensearch.neuralsearch.stats.events.EventStatsManager;
+import org.opensearch.neuralsearch.stats.state.StateStatsManager;
 import org.opensearch.transport.client.Client;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.node.DiscoveryNodes;
@@ -125,7 +127,7 @@ public Collection<Object> createComponents(
         PipelineServiceUtil.instance().initialize(clusterService);
         NeuralSearchSettingsAccessor.instance().initialize(clusterService, environment.settings());
 
-        return List.of(clientAccessor);
+        return List.of(clientAccessor, EventStatsManager.instance(), StateStatsManager.instance());
     }
 
     @Override

src/main/java/org/opensearch/neuralsearch/stats/events/EventStatName.java

@@ -17,10 +17,7 @@
  */
 @Getter
 public enum EventStatName implements StatName {
-    TEXT_EMBEDDING_PROCESSOR_EXECUTIONS("text_embedding_executions", "processors.ingest", EventStatType.TIMESTAMPED_COUNTER),
-
-    // TODO : NEURAL QUERY COUNTER NOT YET IMPLEMENTED
-    NEURAL_QUERY_COUNT("neural_query_count", "query", EventStatType.TIMESTAMPED_COUNTER);
+    TEXT_EMBEDDING_PROCESSOR_EXECUTIONS("text_embedding_executions", "processors.ingest", EventStatType.TIMESTAMPED_COUNTER);
 
     private final String name;
     private final String path;

src/main/java/org/opensearch/neuralsearch/transport/NeuralStatsResponse.java

@@ -4,6 +4,7 @@
  */
 package org.opensearch.neuralsearch.transport;
 
+import lombok.Getter;
 import org.opensearch.action.FailedNodeException;
 import org.opensearch.action.support.nodes.BaseNodesResponse;
 import org.opensearch.cluster.ClusterName;
@@ -22,6 +23,7 @@
 /**
  * NeuralStatsResponse consists of the aggregated responses from the nodes
  */
+@Getter
 public class NeuralStatsResponse extends BaseNodesResponse<NeuralStatsNodeResponse> implements ToXContentObject {
 
     private static final String NODES_KEY = "nodes";
@@ -116,6 +118,7 @@ private Map<String, Object> formatStats(Map<String, StatSnapshot<?>> rawStats) {
 
     private Map<String, Object> getFlattenedStats(Map<String, StatSnapshot<?>> rawStats) {
         if (includeMetadata) {
+            // Safe downcast to object
             return (Map<String, Object>) (Map) rawStats;
         }
         return rawStats.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getValue()));

src/main/java/org/opensearch/neuralsearch/transport/NeuralStatsTransportAction.java

@@ -54,7 +54,9 @@ public NeuralStatsTransportAction(
         ThreadPool threadPool,
         ClusterService clusterService,
         TransportService transportService,
-        ActionFilters actionFilters
+        ActionFilters actionFilters,
+        EventStatsManager eventStatsManager,
+        StateStatsManager stateStatsManager
     ) {
         super(
             NeuralStatsAction.NAME,
@@ -67,8 +69,8 @@ public NeuralStatsTransportAction(
             ThreadPool.Names.MANAGEMENT,
             NeuralStatsNodeResponse.class
         );
-        this.eventStatsManager = EventStatsManager.instance();
-        this.stateStatsManager = StateStatsManager.instance();
+        this.eventStatsManager = eventStatsManager;
+        this.stateStatsManager = stateStatsManager;
     }
 
     @Override

src/test/java/org/opensearch/neuralsearch/rest/RestNeuralStatsHandlerIT.java

@@ -90,6 +90,7 @@ public void test_textEmbedding() throws Exception {
         updateClusterSettings("plugins.neural_search.stats_enabled", false);
         updateClusterSettings("plugins.neural_search.stats_enabled", true);
 
+        // State stats should persist, event stats should be reset
         response = executeNeuralStatRequest(new ArrayList<>(), new ArrayList<>());
         responseBody = EntityUtils.toString(response.getEntity());
         stats = parseStatsResponse(responseBody);
@@ -135,6 +136,7 @@ public void test_includeMetadata() throws Exception {
         Object clusterVersionStatMetadata = getNestedValue(stats, StateStatName.CLUSTER_VERSION.getFullPath());
 
         // Path value should be JSON object (convertible to map)
+        // If metadata wasn't included, path value would be the raw value
         assertTrue(getNestedValue(stats, StateStatName.CLUSTER_VERSION.getFullPath()) instanceof Map);
 
         String statType = ((Map<String, String>) clusterVersionStatMetadata).get(StatSnapshot.STAT_TYPE_FIELD);

src/test/java/org/opensearch/neuralsearch/stats/NeuralStatsInputTests.java

@@ -15,15 +15,18 @@
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Set;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.neuralsearch.util.TestUtils.xContentBuilderToMap;
 
 public class NeuralStatsInputTests extends OpenSearchTestCase {
     private static final String NODE_ID_1 = "node1";
@@ -127,13 +130,13 @@ public void test_toXContent() throws IOException {
 
         XContentBuilder builder = XContentFactory.jsonBuilder();
         input.toXContent(builder, ToXContent.EMPTY_PARAMS);
-        String jsonString = builder.toString();
+        Map<String, Object> responseMap = xContentBuilderToMap(builder);
 
-        assertTrue(jsonString.contains(String.format("\"node_ids\":[\"%s\"]", NODE_ID_1)));
-        assertTrue(jsonString.contains(String.format("\"event_stats\":[\"%s\"]", EVENT_STAT.getName())));
-        assertTrue(jsonString.contains(String.format("\"state_stats\":[\"%s\"]", STATE_STAT.getName())));
-        assertTrue(jsonString.contains("\"include_metadata\":true"));
-        assertTrue(jsonString.contains("\"flat_keys\":true"));
+        assertEquals(Collections.singletonList(NODE_ID_1), responseMap.get("node_ids"));
+        assertEquals(Collections.singletonList(EVENT_STAT.getName()), responseMap.get("event_stats"));
+        assertEquals(Collections.singletonList(STATE_STAT.getName()), responseMap.get("state_stats"));
+        assertEquals(true, responseMap.get("include_metadata"));
+        assertEquals(true, responseMap.get("flat_keys"));
     }
 
     public void test_writeToHandlesEmptyCollections() throws IOException {

src/test/java/org/opensearch/neuralsearch/stats/events/EventStatsManagerTests.java

@@ -96,17 +96,14 @@ public void test_getEventStatSnapshotsReturnsEmptyMap() {
     public void test_reset() {
         // Create multiple mock stats
         TimestampedEventStat mockStat1 = mock(TimestampedEventStat.class);
-        TimestampedEventStat mockStat2 = mock(TimestampedEventStat.class);
 
         // Add mock stats to manager
         eventStatsManager.getStats().clear();
         eventStatsManager.getStats().put(STAT_NAME, mockStat1);
-        eventStatsManager.getStats().put(EventStatName.NEURAL_QUERY_COUNT, mockStat2);
 
         eventStatsManager.reset();
 
         verify(mockStat1, times(1)).reset();
-        verify(mockStat2, times(1)).reset();
     }
 
     public void test_singletonInstanceCreation() {

src/test/java/org/opensearch/neuralsearch/stats/events/TimestampedEventStatSnapshotTests.java

@@ -13,11 +13,13 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Map;
 
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.neuralsearch.util.TestUtils.xContentBuilderToMap;
 
 public class TimestampedEventStatSnapshotTests extends OpenSearchTestCase {
     private static final EventStatName STAT_NAME = EventStatName.TEXT_EMBEDDING_PROCESSOR_EXECUTIONS;
@@ -79,7 +81,7 @@ public void test_aggregateEventStatSnapshotsReturnsNull() {
 
     public void test_aggregateEventStatDataThrowsException() {
         TimestampedEventStatSnapshot snapshot1 = new TimestampedEventStatSnapshot(STAT_NAME, 100L, 50L, 10L);
-        TimestampedEventStatSnapshot snapshot2 = new TimestampedEventStatSnapshot(EventStatName.NEURAL_QUERY_COUNT, 200L, 100L, 5L);
+        TimestampedEventStatSnapshot snapshot2 = new TimestampedEventStatSnapshot(null, 200L, 100L, 5L);
 
         assertThrows(
             IllegalArgumentException.class,
@@ -89,20 +91,15 @@ public void test_aggregateEventStatDataThrowsException() {
 
     public void test_toXContent() throws IOException {
         XContentBuilder builder = JsonXContent.contentBuilder();
-
-        // Create a real snapshot with known values
         TimestampedEventStatSnapshot snapshot = new TimestampedEventStatSnapshot(STAT_NAME, 100L, 50L, 10L);
 
-        // Convert to XContent
         snapshot.toXContent(builder, null);
 
-        // Convert to string for assertion
-        String json = builder.toString();
+        Map<String, Object> responseMap = xContentBuilderToMap(builder);
 
-        // Verify the JSON structure contains all expected fields with correct values
-        assertTrue(json.contains("\"value\":100"));
-        assertTrue(json.contains("\"trailing_interval_value\":50"));
-        assertTrue(json.contains("\"minutes_since_last_event\":10"));
-        assertTrue(json.contains("\"stat_type\""));
+        assertEquals(100, responseMap.get("value"));
+        assertEquals(50, responseMap.get("trailing_interval_value"));
+        assertEquals(10, responseMap.get("minutes_since_last_event"));
+        assertEquals(STAT_NAME.getStatType().getName(), responseMap.get("stat_type"));
     }
 }

src/test/java/org/opensearch/neuralsearch/stats/state/CountableStateStatSnapshotTests.java

@@ -11,6 +11,9 @@
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import java.util.Map;
+
+import static org.opensearch.neuralsearch.util.TestUtils.xContentBuilderToMap;
 
 public class CountableStateStatSnapshotTests extends OpenSearchTestCase {
     private static final StateStatName STAT_NAME = StateStatName.TEXT_EMBEDDING_PROCESSORS;
@@ -26,17 +29,14 @@ public void test_increment() {
 
     public void test_toXContent() throws IOException {
         CountableStateStatSnapshot snapshot = new CountableStateStatSnapshot(STAT_NAME);
-        snapshot.incrementBy(42L);
+        snapshot.incrementBy(8675309L);
 
         XContentBuilder builder = JsonXContent.contentBuilder();
         snapshot.toXContent(builder, ToXContent.EMPTY_PARAMS);
-        String jsonStr = builder.toString();
 
-        // Verify JSON
-        String expectedValueJson = String.format("\"%s\":42", StatSnapshot.VALUE_FIELD);
-        String expectedTypeJson = String.format("\"%s\":\"%s\"", StatSnapshot.STAT_TYPE_FIELD, STAT_NAME.getStatType().getName());
+        Map<String, Object> responseMap = xContentBuilderToMap(builder);
 
-        assertTrue(jsonStr.contains(expectedValueJson));
-        assertTrue(jsonStr.contains(expectedTypeJson));
+        assertEquals(8675309, responseMap.get(StatSnapshot.VALUE_FIELD));
+        assertEquals(STAT_NAME.getStatType().getName(), responseMap.get(StatSnapshot.STAT_TYPE_FIELD));
     }
 }

src/test/java/org/opensearch/neuralsearch/stats/state/SettableStateStatSnapshotTests.java

@@ -11,6 +11,9 @@
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.io.IOException;
+import java.util.Map;
+
+import static org.opensearch.neuralsearch.util.TestUtils.xContentBuilderToMap;
 
 public class SettableStateStatSnapshotTests extends OpenSearchTestCase {
 
@@ -35,17 +38,12 @@ public void test_setValueUpdates() {
 
     public void test_toXContent() throws IOException {
         SettableStateStatSnapshot<String> snapshot = new SettableStateStatSnapshot<>(STAT_NAME, SETTABLE_VALUE);
-
         XContentBuilder builder = JsonXContent.contentBuilder();
-
         snapshot.toXContent(builder, ToXContent.EMPTY_PARAMS);
-        String json = builder.toString();
 
-        // Verify the JSON structure
-        String expectedValueJson = String.format("\"%s\":\"%s\"", StatSnapshot.VALUE_FIELD, SETTABLE_VALUE);
-        String expectedTypeJson = String.format("\"%s\":\"%s\"", StatSnapshot.STAT_TYPE_FIELD, STAT_NAME.getStatType().getName());
+        Map<String, Object> responseMap = xContentBuilderToMap(builder);
 
-        assertTrue(json.contains(expectedValueJson));
-        assertTrue(json.contains(expectedTypeJson));
+        assertEquals(SETTABLE_VALUE, responseMap.get(StatSnapshot.VALUE_FIELD));
+        assertEquals(STAT_NAME.getStatType().getName(), responseMap.get(StatSnapshot.STAT_TYPE_FIELD));
     }
 }