-

internal/proxy/proxy_test.go

@@ -0,0 +1,65 @@
+package proxy
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strings"
+	"testing"
+
+	"github.com/example/ollama-fair-gateway/internal/cost"
+)
+
+func TestNativeMeter(t *testing.T) {
+	m := newMeter("ollama", 0)
+	m.Feed([]byte("{\"response\":\"x\",\"done\":false}\n{\"done\":true,\"prompt_eval_count\":12,\"eval_count\":4,\"eval_duration\":100}\n"))
+	u := m.Finish(100)
+	if u.PromptTokens != 12 || u.CompletionTokens != 4 || u.EvalNS != 100 || u.Approximate {
+		t.Fatalf("bad usage %#v", u)
+	}
+}
+func TestOpenAIMeter(t *testing.T) {
+	m := newMeter("openai", 9)
+	m.Feed([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\ndata: {\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2}}\n"))
+	u := m.Finish(100)
+	if u.PromptTokens != 10 || u.CompletionTokens != 2 || u.Approximate {
+		t.Fatalf("bad usage %#v", u)
+	}
+}
+
+func TestAnthropicMeter(t *testing.T) {
+	m := newMeter("anthropic", 0)
+	m.Feed([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":11,\"output_tokens\":0}}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":5}}\n\n"))
+	u := m.Finish(100)
+	if u.PromptTokens != 11 || u.CompletionTokens != 5 || u.Approximate {
+		t.Fatalf("bad anthropic usage %#v", u)
+	}
+}
+
+func TestForwardProgressObserver(t *testing.T) {
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/x-ndjson")
+		_, _ = io.WriteString(w, "{\"message\":{\"content\":\"hi\"},\"done\":false}\n")
+		if f, ok := w.(http.Flusher); ok {
+			f.Flush()
+		}
+		_, _ = io.WriteString(w, "{\"done\":true,\"prompt_eval_count\":10,\"eval_count\":2}\n")
+	}))
+	defer backend.Close()
+	target, _ := url.Parse(backend.URL)
+	in := httptest.NewRequest(http.MethodPost, "http://gateway/api/chat", strings.NewReader(`{"model":"x"}`))
+	out := httptest.NewRecorder()
+	var calls int
+	var lastBytes int64
+	var last cost.Usage
+	res := New().Forward(context.Background(), out, in, target, strings.NewReader(`{"model":"x"}`), "ollama", 0, func(n int64, u cost.Usage) {
+		calls++
+		lastBytes = n
+		last = u
+	})
+	if res.Status != http.StatusOK || calls == 0 || lastBytes == 0 || last.PromptTokens != 10 || last.CompletionTokens != 2 {
+		t.Fatalf("status=%d calls=%d bytes=%d usage=%#v", res.Status, calls, lastBytes, last)
+	}
+}