[Java] Add test utility for stubbing addition of counters.

G-Research · Dec 17, 2024 · fc97fba · fc97fba
1 parent 13fd10c
commit fc97fba
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 12 deletions.
diff --git a/aeron-cluster/src/test/java/io/aeron/cluster/ConsensusModuleContextTest.java b/aeron-cluster/src/test/java/io/aeron/cluster/ConsensusModuleContextTest.java
@@ -91,18 +91,7 @@ void beforeEach()
         final Aeron aeron = mock(Aeron.class);
         when(aeron.addCounter(
             anyInt(), any(DirectBuffer.class), anyInt(), anyInt(), any(DirectBuffer.class), anyInt(), anyInt()))
-            .thenAnswer(invocation ->
-            {
-                final int counterId = countersManager.allocate(
-                    invocation.getArgument(0),
-                    invocation.getArgument(1),
-                    invocation.getArgument(2),
-                    invocation.getArgument(3),
-                    invocation.getArgument(4),
-                    invocation.getArgument(5),
-                    invocation.getArgument(6));
-                return new Counter(countersManager, registrationId++, counterId);
-            });
+            .thenAnswer(Tests.addCounterAnswer(countersManager, () -> registrationId++));
         when(aeron.context()).thenReturn(aeronContext);
         when(aeron.conductorAgentInvoker()).thenReturn(conductorInvoker);
         when(aeron.countersReader()).thenReturn(countersManager);

diff --git a/aeron-test-support/src/main/java/io/aeron/test/Tests.java b/aeron-test-support/src/main/java/io/aeron/test/Tests.java
@@ -17,6 +17,7 @@
 package io.aeron.test;
 
 import io.aeron.Aeron;
+import io.aeron.Counter;
 import io.aeron.Publication;
 import io.aeron.Subscription;
 import io.aeron.archive.status.RecordingPos;
@@ -25,6 +26,7 @@
 import io.aeron.exceptions.RegistrationException;
 import io.aeron.exceptions.TimeoutException;
 
+import org.agrona.DirectBuffer;
 import org.agrona.LangUtil;
 import org.agrona.concurrent.IdleStrategy;
 import org.agrona.concurrent.SleepingMillisIdleStrategy;
@@ -35,6 +37,7 @@
 import org.agrona.concurrent.status.CountersReader;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.TestWatcher;
+import org.mockito.stubbing.Answer;
 
 import javax.management.InstanceNotFoundException;
 import javax.management.MBeanServer;
@@ -832,6 +835,27 @@ public static CountersManager newCountersManager(final int dataLength)
             new UnsafeBuffer(ByteBuffer.allocateDirect(dataLength)));
     }
 
+    public static Answer<Counter> addCounterAnswer(
+        final CountersManager countersManager,
+        final LongSupplier registrationId)
+    {
+        return invocation ->
+        {
+            final int counterType = invocation.getArgument(0, Integer.class);
+            final DirectBuffer keyBuffer = invocation.getArgument(1, DirectBuffer.class);
+            final int keyOffset = invocation.getArgument(2, Integer.class);
+            final int keyLength = invocation.getArgument(3, Integer.class);
+            final DirectBuffer labelBuffer = invocation.getArgument(4, DirectBuffer.class);
+            final int labelOffset = invocation.getArgument(5, Integer.class);
+            final int labelLength = invocation.getArgument(6, Integer.class);
+
+            final int allocate = countersManager.allocate(
+                counterType, keyBuffer, keyOffset, keyLength, labelBuffer, labelOffset, labelLength);
+
+            return new Counter(countersManager, registrationId.getAsLong(), allocate);
+        };
+    }
+
     public static Throwable setOrUpdateError(final Throwable existingError, final Throwable newError)
     {
         if (null == existingError)