ai-gateway-go/internal/trace/store_integration_test.go

package trace

import (
	"context"
	"os"
	"strings"
	"testing"
	"time"

	"aigateway.local/core/internal/platform/config"
	"aigateway.local/core/internal/platform/database"
)

func TestTracePostgreSQLLifecycle(t *testing.T) {
	databaseURL := os.Getenv("TRACE_TEST_DATABASE_URL")
	if databaseURL == "" {
		t.Skip("TRACE_TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := database.Open(ctx, config.Database{URL: databaseURL, MaxConns: 4, MinConns: 0})
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	store := NewStore(pool)
	_, _ = pool.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE request_id LIKE 'trace-integration%'`)
	item, err := store.Start(ctx, StartInput{RequestID: "trace-integration", TraceType: "application", TargetID: "77777777-7777-4777-8777-777777777777", TargetCode: "trace_app", ConversationID: "conversation-1", Metadata: map[string]any{"version": 1}})
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Exec(ctx, `DELETE FROM gateway.agent_traces WHERE request_id LIKE 'trace-integration%'`)
	model, err := store.StartSpan(ctx, SpanInput{TraceID: item.ID, SpanType: "model", Name: "chat.completions", Round: 0, ProviderCode: "test", Model: "test-model"})
	if err != nil {
		t.Fatal(err)
	}
	if err = store.FinishSpan(ctx, model.ID, SpanFinishInput{Status: "success", InputTokens: 12, OutputTokens: 8, Metadata: map[string]any{"http_status": 200}}); err != nil {
		t.Fatal(err)
	}
	tool, err := store.StartSpan(ctx, SpanInput{TraceID: item.ID, SpanType: "tool", Name: "lookup", Round: 0})
	if err != nil {
		t.Fatal(err)
	}
	if err = store.FinishSpan(ctx, tool.ID, SpanFinishInput{Status: "error", Error: "upstream timeout"}); err != nil {
		t.Fatal(err)
	}
	if err = store.Finish(ctx, item.ID, FinishInput{Status: "success", RetrievalCount: 2, ModelCallCount: 1, ToolCallCount: 1}); err != nil {
		t.Fatal(err)
	}
	loaded, err := store.Get(ctx, item.ID)
	if err != nil || loaded.Status != "success" || loaded.ModelCallCount != 1 || len(loaded.Spans) != 2 {
		t.Fatalf("loaded=%+v err=%v", loaded, err)
	}
	if loaded.Spans[0].ProviderCode != "test" || loaded.Spans[0].Model != "test-model" {
		t.Fatalf("model span route metadata was not preserved: %+v", loaded.Spans[0])
	}
	if loaded.LatencyMS == nil || loaded.StartedAt.After(time.Now().UTC().Add(time.Second)) {
		t.Fatalf("invalid timing: %+v", loaded)
	}
	items, err := store.List(ctx, Filter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), TargetCode: "trace_app", Limit: 10})
	if err != nil || len(items) != 1 || items[0].ID != item.ID {
		t.Fatalf("list=%+v err=%v", items, err)
	}
	sessions, err := store.ListSessions(ctx, SessionFilter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), SessionID: "application:trace_app:conversation-1", Limit: 10})
	if err != nil || len(sessions) != 1 || sessions[0].TraceCount != 1 || sessions[0].LatestTraceID != item.ID || sessions[0].ModelCallCount != 1 {
		t.Fatalf("sessions=%+v err=%v", sessions, err)
	}
	digital, err := store.Start(ctx, StartInput{RequestID: "trace-integration-digital", TraceType: "digital_employee", TargetID: "88888888-8888-4888-8888-888888888888", TargetCode: "trace_employee", ConversationID: "conversation-1"})
	if err != nil {
		t.Fatal(err)
	}
	if err = store.Finish(ctx, digital.ID, FinishInput{Status: "success", ModelCallCount: 1}); err != nil {
		t.Fatal(err)
	}
	allSessions, err := store.ListSessions(ctx, SessionFilter{From: time.Now().Add(-time.Minute), To: time.Now().Add(time.Minute), Limit: 10})
	if err != nil || len(allSessions) < 2 {
		t.Fatalf("all sessions=%+v err=%v", allSessions, err)
	}
	seenTypes := map[string]bool{}
	for _, session := range allSessions {
		if strings.HasPrefix(session.ID, "application:trace_app:") || strings.HasPrefix(session.ID, "digital_employee:trace_employee:") {
			seenTypes[session.TraceType] = true
		}
	}
	if !seenTypes["application"] || !seenTypes["digital_employee"] {
		t.Fatalf("session types=%v", seenTypes)
	}
}