DataStreamMap assertion in test

ivandika3 · Dec 24, 2024 · 9f05a09 · 9f05a09
1 parent 13b8cdd
commit 9f05a09
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/ratis-server-api/src/main/java/org/apache/ratis/server/DataStreamMap.java b/ratis-server-api/src/main/java/org/apache/ratis/server/DataStreamMap.java
@@ -31,4 +31,7 @@ CompletableFuture<DataStream> computeIfAbsent(ClientInvocationId invocationId,
 
   /** Similar to {@link java.util.Map#remove(java.lang.Object). */
   CompletableFuture<DataStream> remove(ClientInvocationId invocationId);
+
+  /** Similar to {@link java.util.Map#size()}. */
+  int size();
 }
diff --git a/ratis-server/src/main/java/org/apache/ratis/server/impl/DataStreamMapImpl.java b/ratis-server/src/main/java/org/apache/ratis/server/impl/DataStreamMapImpl.java
@@ -50,6 +50,11 @@ public CompletableFuture<DataStream> computeIfAbsent(ClientInvocationId invocati
     return map.computeIfAbsent(invocationId, newDataStream);
   }
 
+  @Override
+  public int size() {
+    return map.size();
+  }
+
   @Override
   public String toString() {
     return name;

diff --git a/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamAsyncClusterTests.java b/ratis-test/src/test/java/org/apache/ratis/datastream/DataStreamAsyncClusterTests.java
@@ -143,6 +143,7 @@ void runTestDataStream(CLUSTER cluster, boolean stepDownLeader,
         Assertions.assertFalse(s.getDataChannel().isOpen());
         DataStreamTestUtils.assertLogEntry(impl, s);
       }
+      Assertions.assertEquals(0, impl.getDataStreamMap().size());
     }
   }
-Original file line number
+Diff line change
@@ Expand Up @@
             Assertions.assertFalse(s.getDataChannel().isOpen());
             DataStreamTestUtils.assertLogEntry(impl, s);
           }
+          Assertions.assertEquals(0, impl.getDataStreamMap().size());
         }
       }
@@ Expand Down @@