Merge pull request #25 from braintrustdata/ark/vcr-wiremock

VCR util for unit tests

Author: realark

scripts/erase-cassettes.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+cd "$(dirname "$(readlink -f "${BASH_SOURCE}")")"/..
+
+rm -rf src/test/resources/cassettes/
+mkdir -p src/test/resources/cassettes/

scripts/record-cassettes.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+cd "$(dirname "$(readlink -f "${BASH_SOURCE}")")"/..
+
+export VCR_MODE=record
+./gradlew test

src/test/java/dev/braintrust/TestHarness.java

@@ -1,7 +1,6 @@
 package dev.braintrust;
 
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.fail;
 
 import dev.braintrust.api.BraintrustApiClient;
 import dev.braintrust.config.BraintrustConfig;
@@ -14,7 +13,6 @@
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.logs.SdkLoggerProvider;
 import io.opentelemetry.sdk.metrics.SdkMeterProvider;
-import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
@@ -27,6 +25,18 @@
 import lombok.experimental.Accessors;
 
 public class TestHarness {
+    private static final VCR vcr;
+
+    static {
+        vcr =
+                new VCR(
+                        java.util.Map.of(
+                                "https://api.openai.com/v1", "openai",
+                                "https://api.anthropic.com", "anthropic",
+                                "https://generativelanguage.googleapis.com", "google"));
+        vcr.start();
+        Runtime.getRuntime().addShutdownHook(new Thread(vcr::stop));
+    }
 
     public static TestHarness setup() {
         return setup(createTestConfig());
@@ -75,13 +85,12 @@ public static synchronized TestHarness setup(BraintrustConfig config) {
     @Accessors(fluent = true)
     private final Braintrust braintrust;
 
-    private final @Nonnull InMemorySpanExporter spanExporter;
+    private final @Nonnull UnitTestSpanExporter spanExporter;
 
     private TestHarness(@Nonnull Braintrust braintrust) {
         this.braintrust = braintrust;
-
         var tracerBuilder = SdkTracerProvider.builder();
-        this.spanExporter = InMemorySpanExporter.create();
+        this.spanExporter = new UnitTestSpanExporter();
         var loggerBuilder = SdkLoggerProvider.builder();
         var meterBuilder = SdkMeterProvider.builder();
         braintrust.openTelemetryEnable(tracerBuilder, loggerBuilder, meterBuilder);
@@ -102,6 +111,30 @@ private TestHarness(@Nonnull Braintrust braintrust) {
         this.openTelemetry = openTelemetry;
     }
 
+    public String openAiBaseUrl() {
+        return vcr.getUrlForTargetBase("https://api.openai.com/v1");
+    }
+
+    public String openAiApiKey() {
+        return getEnv("OPENAI_API_KEY", "test-key");
+    }
+
+    public String anthropicBaseUrl() {
+        return vcr.getUrlForTargetBase("https://api.anthropic.com");
+    }
+
+    public String anthropicApiKey() {
+        return getEnv("ANTHROPIC_API_KEY", "test-key");
+    }
+
+    public String googleBaseUrl() {
+        return vcr.getUrlForTargetBase("https://generativelanguage.googleapis.com");
+    }
+
+    public String googleApiKey() {
+        return getEnv("GOOGLE_API_KEY", getEnv("GEMINI_API_KEY", "test-key"));
+    }
+
     /** flush all pending spans and return all spans which have been exported so far */
     public List<SpanData> awaitExportedSpans() {
         assertTrue(
@@ -120,21 +153,7 @@ public List<SpanData> awaitExportedSpans() {
      */
     @SneakyThrows
     public List<SpanData> awaitExportedSpans(int minSpanCount) {
-        var spans = awaitExportedSpans();
-        int attempts = 0;
-        while (spans.size() < minSpanCount) {
-            attempts++;
-            if (attempts > 30) {
-                fail(
-                        String.format(
-                                "Timeout waiting for spans: expected at least %d spans, but got %d"
-                                        + " after %d attempts",
-                                minSpanCount, spans.size(), attempts));
-            }
-            Thread.sleep(1000);
-            spans = awaitExportedSpans();
-        }
-        return spans;
+        return spanExporter.getFinishedSpanItems(minSpanCount);
     }
 
     private static BraintrustApiClient.InMemoryImpl createApiClient() {
@@ -163,4 +182,9 @@ public static BraintrustConfig createTestConfig() {
                 "BRAINTRUST_APP_URL", "https://testhost:3000",
                 "BRAINTRUST_DEFAULT_PROJECT_NAME", defaultProjectName());
     }
+
+    private static String getEnv(String envarName, String defaultValue) {
+        var envar = System.getenv(envarName);
+        return envar == null ? defaultValue : envar;
+    }
 }

src/test/java/dev/braintrust/UnitTestSpanExporter.java

@@ -0,0 +1,84 @@
+package dev.braintrust;
+
+import static org.junit.jupiter.api.Assertions.fail;
+
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.trace.data.SpanData;
+import io.opentelemetry.sdk.trace.export.SpanExporter;
+import java.util.*;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import lombok.SneakyThrows;
+
+/**
+ * an in memory span exporter which allows for blocking until the exported span count reaches an
+ * expected min
+ */
+public class UnitTestSpanExporter implements SpanExporter {
+    private final Queue<SpanData> finishedSpanItems = new ConcurrentLinkedQueue<>();
+    private final Lock lock = new ReentrantLock();
+    private final Condition spansAdded = lock.newCondition();
+    private boolean isStopped = false;
+
+    public UnitTestSpanExporter() {}
+
+    @SneakyThrows
+    public List<SpanData> getFinishedSpanItems(int minSpanCount) {
+        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30);
+        lock.lock();
+        try {
+            while (finishedSpanItems.size() < minSpanCount) {
+                long remainingNanos = deadline - System.nanoTime();
+                if (remainingNanos <= 0) {
+                    fail(
+                            String.format(
+                                    "Timeout waiting for spans: expected at least %d spans, but got"
+                                            + " %d after 30 seconds",
+                                    minSpanCount, finishedSpanItems.size()));
+                }
+                spansAdded.awaitNanos(remainingNanos);
+            }
+            return Collections.unmodifiableList(new ArrayList<>(finishedSpanItems));
+        } finally {
+            lock.unlock();
+        }
+    }
+
+    public List<SpanData> getFinishedSpanItems() {
+        return Collections.unmodifiableList(new ArrayList<>(finishedSpanItems));
+    }
+
+    public void reset() {
+        finishedSpanItems.clear();
+    }
+
+    @Override
+    public CompletableResultCode export(Collection<SpanData> spans) {
+        if (isStopped) {
+            return CompletableResultCode.ofFailure();
+        }
+        lock.lock();
+        try {
+            finishedSpanItems.addAll(spans);
+            spansAdded.signalAll();
+        } finally {
+            lock.unlock();
+        }
+        return CompletableResultCode.ofSuccess();
+    }
+
+    @Override
+    public CompletableResultCode flush() {
+        return CompletableResultCode.ofSuccess();
+    }
+
+    @Override
+    public CompletableResultCode shutdown() {
+        finishedSpanItems.clear();
+        isStopped = true;
+        return CompletableResultCode.ofSuccess();
+    }
+}