Vectrum
Sign In
148 lines · 3.8 KB · Go
package metrics

import (
	"bytes"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

func decode(t *testing.T, line string) map[string]any {
	t.Helper()
	var m map[string]any
	if err := json.Unmarshal([]byte(line), &m); err != nil {
		t.Fatalf("line is not valid JSON (%v): %s", err, line)
	}
	return m
}

func TestLogWritesOneObjectPerLine(t *testing.T) {
	var buf bytes.Buffer
	l := NewWriter(&buf, "sess-1")
	l.Log(Event{Type: "turn", Turn: 1, InputTokens: 900})
	l.Log(Event{Type: "turn", Turn: 2, InputTokens: 1200})

	lines := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
	if len(lines) != 2 {
		t.Fatalf("got %d lines, want 2:\n%s", len(lines), buf.String())
	}
	first := decode(t, lines[0])
	if first["session"] != "sess-1" {
		t.Errorf("session = %v, want sess-1 stamped automatically", first["session"])
	}
	if first["ts"] == nil {
		t.Error("ts must be filled in when the caller leaves it zero")
	}
	if first["input_tokens"].(float64) != 900 {
		t.Errorf("input_tokens = %v, want 900", first["input_tokens"])
	}
}

func TestExtraIsFlattened(t *testing.T) {
	var buf bytes.Buffer
	NewWriter(&buf, "s").Log(Event{
		Type:  "edit",
		Tool:  "edit_file",
		Extra: map[string]any{"matches": 3, "path": "main.go"},
	})
	m := decode(t, buf.String())
	if m["matches"].(float64) != 3 || m["path"] != "main.go" {
		t.Errorf("Extra should sit alongside declared fields, got %v", m)
	}
}

func TestExtraCannotClobberDeclaredFields(t *testing.T) {
	var buf bytes.Buffer
	NewWriter(&buf, "s").Log(Event{Type: "edit", Extra: map[string]any{"type": "spoofed"}})
	if m := decode(t, buf.String()); m["type"] != "edit" {
		t.Errorf("type = %v, want the declared value to win", m["type"])
	}
}

func TestToolCallRecordsOutcome(t *testing.T) {
	tests := []struct {
		name   string
		err    error
		wantOK bool
		wantEr string
	}{
		{name: "success", err: nil, wantOK: true},
		{name: "failure", err: errors.New("no such file"), wantOK: false, wantEr: "no such file"},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var buf bytes.Buffer
			NewWriter(&buf, "s").ToolCall("read_file", 12*time.Millisecond, tc.err, nil)
			m := decode(t, buf.String())
			if m["type"] != "tool_call" || m["tool"] != "read_file" {
				t.Fatalf("unexpected event %v", m)
			}
			if m["ok"] != tc.wantOK {
				t.Errorf("ok = %v, want %v", m["ok"], tc.wantOK)
			}
			if m["dur_ms"].(float64) != 12 {
				t.Errorf("dur_ms = %v, want 12", m["dur_ms"])
			}
			if tc.wantEr != "" && m["err"] != tc.wantEr {
				t.Errorf("err = %v, want %q", m["err"], tc.wantEr)
			}
		})
	}
}

func TestNilLoggerIsANoOp(t *testing.T) {
	var l *Logger
	l.Log(Event{Type: "turn"})
	l.ToolCall("x", 0, nil, nil)
	if l.Session() != "" {
		t.Error("nil logger should report an empty session")
	}
	if err := l.Close(); err != nil {
		t.Errorf("Close on nil logger = %v, want nil", err)
	}
}

func TestOpenCreatesParentDirsAndAppends(t *testing.T) {
	path := filepath.Join(t.TempDir(), ".harness", "events.jsonl")
	for i := 0; i < 2; i++ {
		l, err := Open(path, "s")
		if err != nil {
			t.Fatalf("Open: %v", err)
		}
		l.Log(Event{Type: "run_start"})
		l.Close()
	}
	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	if n := strings.Count(string(b), "\n"); n != 2 {
		t.Errorf("got %d lines, want 2 (second Open must append, not truncate)", n)
	}
}

func TestConcurrentLogLinesStayIntact(t *testing.T) {
	var buf bytes.Buffer
	l := NewWriter(&buf, "s")
	var wg sync.WaitGroup
	for i := 0; i < 50; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			l.Log(Event{Type: "tool_call", Turn: i})
		}(i)
	}
	wg.Wait()
	lines := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
	if len(lines) != 50 {
		t.Fatalf("got %d lines, want 50", len(lines))
	}
	for _, ln := range lines {
		decode(t, ln)
	}
}