fix tests

LangStream · Nov 7, 2023 · 90ce550 · 90ce550
1 parent 376e191
commit 90ce550
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/...stream-agents/langstream-agent-s3/src/test/java/ai/langstream/agents/s3/S3SourceTest.java b/...stream-agents/langstream-agent-s3/src/test/java/ai/langstream/agents/s3/S3SourceTest.java
@@ -19,10 +19,14 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 import static org.testcontainers.containers.localstack.LocalStackContainer.Service.S3;
 
 import ai.langstream.api.runner.code.AgentCodeRegistry;
+import ai.langstream.api.runner.code.AgentContext;
 import ai.langstream.api.runner.code.AgentSource;
+import ai.langstream.api.runner.code.MetricsReporter;
 import ai.langstream.api.runner.code.Record;
 import io.minio.ListObjectsArgs;
 import io.minio.MinioClient;
@@ -175,6 +179,9 @@ private AgentSource buildAgentSource(String bucket) throws Exception {
         configs.put("endpoint", endpoint);
         configs.put("bucketName", bucket);
         agentSource.init(configs);
+        AgentContext context = mock(AgentContext.class);
+        when(context.getMetricsReporter()).thenReturn(MetricsReporter.DISABLED);
+        agentSource.setContext(context);
         agentSource.start();
         return agentSource;
     }

diff --git a/...runtime/langstream-runtime-impl/src/test/java/ai/langstream/kafka/GenIAgentsRunnerIT.java b/...runtime/langstream-runtime-impl/src/test/java/ai/langstream/kafka/GenIAgentsRunnerIT.java
@@ -17,6 +17,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import ai.langstream.api.runner.code.AgentStatusResponse;
 import java.nio.charset.StandardCharsets;
@@ -150,8 +151,8 @@ public void testRunAIToolsComposite() throws Exception {
                     default:
                         throw new IllegalStateException("Unexpected value: " + p.getAgentId());
                 }
-                assertEquals(0L, p.getMetrics().getTotalIn());
-                assertEquals(0L, p.getMetrics().getTotalOut());
+                assertTrue(p.getMetrics().getTotalIn() > 0);
+                assertTrue(p.getMetrics().getTotalOut() > 0);
                 if (!mayHaveProcessed) {
                     assertEquals(0L, p.getMetrics().getLastProcessedAt());
                 }

diff --git a/...untime/langstream-runtime-impl/src/test/java/ai/langstream/runtime/LoadAgentCodeTest.java b/...untime/langstream-runtime-impl/src/test/java/ai/langstream/runtime/LoadAgentCodeTest.java
@@ -18,10 +18,14 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import ai.langstream.api.runner.code.AgentCodeRegistry;
+import ai.langstream.api.runner.code.AgentContext;
 import ai.langstream.api.runner.code.AgentProcessor;
 import ai.langstream.api.runner.code.Header;
+import ai.langstream.api.runner.code.MetricsReporter;
 import ai.langstream.api.runner.code.Record;
 import java.util.ArrayList;
 import java.util.List;
@@ -35,6 +39,7 @@ class LoadAgentCodeTest {
     public void testLoadNoop() throws Exception {
         AgentCodeRegistry registry = new AgentCodeRegistry();
         AgentProcessor noop = (AgentProcessor) registry.getAgentCode("noop").agentCode();
+        setContext(noop);
         MyRecord myRecord = new MyRecord();
         List<AgentProcessor.SourceRecordAndResult> res = new ArrayList<>();
         noop.process(List.of(myRecord), res::add);
@@ -46,6 +51,7 @@ public void testLoadNoop() throws Exception {
     public void testLoadIdentity() throws Exception {
         AgentCodeRegistry registry = new AgentCodeRegistry();
         AgentProcessor noop = (AgentProcessor) registry.getAgentCode("identity").agentCode();
+        setContext(noop);
         MyRecord myRecord = new MyRecord();
         List<AgentProcessor.SourceRecordAndResult> res = new ArrayList<>();
         noop.process(List.of(myRecord), res::add);
@@ -58,6 +64,12 @@ public void testLoadIdentity() throws Exception {
         assertSame(myRecord, res.get(0).sourceRecord());
     }
 
+    private static void setContext(AgentProcessor noop) throws Exception {
+        AgentContext agentContext = mock(AgentContext.class);
+        noop.setContext(agentContext);
+        when(agentContext.getMetricsReporter()).thenReturn(MetricsReporter.DISABLED);
+    }
+
     private static class MyRecord implements Record {
         @Override
         public Object key() {