adding boolean check if model stats exists

Signed-off-by: Dhrubo Saha <[email protected]>
opensearch-project · Oct 4, 2023 · 4abc1d1 · 4abc1d1
1 parent f99a063
commit 4abc1d1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/stats/MLModelStats.java b/plugin/src/main/java/org/opensearch/ml/stats/MLModelStats.java
@@ -27,7 +27,9 @@ public class MLModelStats implements ToXContentFragment, Writeable {
     private Map<ActionName, MLActionStats> modelStats;
 
     public MLModelStats(StreamInput in) throws IOException {
-        this.modelStats = in.readMap(stream -> stream.readEnum(ActionName.class), MLActionStats::new);
+        if (in.readBoolean()) {
+            this.modelStats = in.readMap(stream -> stream.readEnum(ActionName.class), MLActionStats::new);
+        }
     }
 
     public MLModelStats(Map<ActionName, MLActionStats> modelStats) {
@@ -37,7 +39,10 @@ public MLModelStats(Map<ActionName, MLActionStats> modelStats) {
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         if (modelStats != null && modelStats.size() > 0) {
+            out.writeBoolean(true);
             out.writeMap(modelStats, (stream, v) -> stream.writeEnum(v), (stream, stats) -> stats.writeTo(stream));
+        } else {
+            out.writeBoolean(false);
         }
     }
 

diff --git a/plugin/src/test/java/org/opensearch/ml/stats/MLModelStatsTests.java b/plugin/src/test/java/org/opensearch/ml/stats/MLModelStatsTests.java
@@ -57,7 +57,10 @@ public void testEmptySerializationDeserialization() throws IOException {
         MLModelStats mlModelEmptyStats = new MLModelStats(modelStats);
         BytesStreamOutput output = new BytesStreamOutput();
         mlModelEmptyStats.writeTo(output);
-        assertEquals(0, output.bytes().length());
+        MLModelStats parsedMLModelStats = new MLModelStats(output.bytes().streamInput());
+        MLActionStats parsedMLActionStats = parsedMLModelStats.getActionStats(ActionName.PREDICT);
+        assertNull(parsedMLActionStats);
+        // assertEquals(0, output.bytes().length());
     }
 
     public void testToXContent() throws IOException {