fathom/internal/store/chunk_test.go

package store

import (
	"math/rand"
	"reflect"
	"testing"

	"git.z.gd/zzg/fathom/internal/model"
)

func TestDictRoundTrip(t *testing.T) {
	d := NewDict()
	a, b, c := d.Put("alpha"), d.Put("beta"), d.Put("alpha")
	if a != c {
		t.Fatalf("interning the same string gave codes %d and %d", a, c)
	}
	if a == b {
		t.Fatalf("distinct strings collided on code %d", a)
	}
	if got := d.Get(b); got != "beta" {
		t.Fatalf("Get(%d) = %q, want %q", b, got, "beta")
	}
	if _, ok := d.Lookup("gamma"); ok {
		t.Fatal("Lookup found a string that was never interned")
	}
	// Out-of-range codes come from damaged chunks and must not panic.
	if got := d.Get(9999); got != "" {
		t.Fatalf("Get(out-of-range) = %q, want empty", got)
	}
}

func sampleSpans() []model.Span {
	return []model.Span{
		{
			TraceID: "trace-1", SpanID: "span-1", ParentSpanID: "",
			Name: "claude_code.interaction", Service: "claude-code",
			Kind: model.KindAgentStep, Status: model.StatusOK,
			StartTS: 1_700_000_000_000_000, EndTS: 1_700_000_015_396_000,
			SampleRate: 1,
			Attrs: model.Attrs{
				"span.type":    model.StringValue("interaction"),
				"input_tokens": model.IntValue(4213),
				"cost_usd":     model.FloatValue(0.0231),
				"success":      model.BoolValue(true),
			},
			Input:  "Use ls to list the .txt files in this directory.",
			Output: ".txt files: a.txt, b.txt, c.txt",
		},
		{
			TraceID: "trace-1", SpanID: "span-2", ParentSpanID: "span-1",
			Name: "claude_code.tool", Service: "claude-code",
			Kind: model.KindToolCall, Status: model.StatusError,
			StartTS: 1_700_000_003_000_000, EndTS: 1_700_000_003_092_000,
			SampleRate: 0.5,
			Attrs: model.Attrs{
				"tool_name":   model.StringValue("Bash"),
				"tool_use_id": model.StringValue("toolu_01TjbRiv"),
			},
		},
		{
			// No attributes, no blobs: the common shape, and the one most
			// likely to expose an off-by-one in the attribute windows.
			TraceID: "trace-2", SpanID: "span-3", ParentSpanID: "",
			Name: "claude_code.llm_request", Service: "claude-code",
			Kind: model.KindLLMCall, Status: model.StatusUnset,
			StartTS: 1_700_000_020_000_000, EndTS: 1_700_000_023_253_000,
			SampleRate: 1,
		},
	}
}

func TestChunkRoundTrip(t *testing.T) {
	blob, err := OpenBlobFile(t.TempDir() + "/x.blob")
	if err != nil {
		t.Fatal(err)
	}
	defer blob.Close()

	spans := sampleSpans()
	chunk, err := buildSpanChunk(spans, blob)
	if err != nil {
		t.Fatal(err)
	}

	var buf bytesBuffer
	if err := chunk.Encode(&buf); err != nil {
		t.Fatal(err)
	}
	got, err := DecodeChunk(buf.b)
	if err != nil {
		t.Fatal(err)
	}

	if got.Rows != len(spans) {
		t.Fatalf("decoded %d rows, wrote %d", got.Rows, len(spans))
	}
	for i, want := range spans {
		g := readSpan(got, i)
		// Blobs are deliberately not part of the row; compare them separately.
		want.Input, want.Output = "", ""
		if len(want.Attrs) == 0 {
			want.Attrs = nil
		}
		if !reflect.DeepEqual(g, want) {
			t.Errorf("row %d round-tripped wrong:\n got %+v\nwant %+v", i, g, want)
		}
	}

	// Blob refs must resolve to the original text.
	for i, want := range spans {
		in, err := blob.Read(blobAt(got, ColInput, i))
		if err != nil {
			t.Fatalf("row %d input: %v", i, err)
		}
		if in != want.Input {
			t.Errorf("row %d input = %q, want %q", i, in, want.Input)
		}
		out, err := blob.Read(blobAt(got, ColOutput, i))
		if err != nil {
			t.Fatalf("row %d output: %v", i, err)
		}
		if out != want.Output {
			t.Errorf("row %d output = %q, want %q", i, out, want.Output)
		}
	}
}

func TestChunkEncodingIsDeterministic(t *testing.T) {
	// Attribute maps iterate in random order in Go. If that order leaked into
	// the encoding, chunks would differ run to run and this test would be flaky
	// rather than failing outright, so it encodes repeatedly.
	dir := t.TempDir()
	var first []byte
	for i := 0; i < 8; i++ {
		blob, err := OpenBlobFile(dir + "/" + string(rune('a'+i)) + ".blob")
		if err != nil {
			t.Fatal(err)
		}
		chunk, err := buildSpanChunk(sampleSpans(), blob)
		if err != nil {
			t.Fatal(err)
		}
		var buf bytesBuffer
		if err := chunk.Encode(&buf); err != nil {
			t.Fatal(err)
		}
		blob.Close()
		if i == 0 {
			first = buf.b
			continue
		}
		if !reflect.DeepEqual(first, buf.b) {
			t.Fatalf("encoding run %d differs from run 0; attribute ordering is leaking", i)
		}
	}
}

func TestDecodeChunkRejectsGarbage(t *testing.T) {
	blob, _ := OpenBlobFile(t.TempDir() + "/x.blob")
	defer blob.Close()
	chunk, err := buildSpanChunk(sampleSpans(), blob)
	if err != nil {
		t.Fatal(err)
	}
	var buf bytesBuffer
	if err := chunk.Encode(&buf); err != nil {
		t.Fatal(err)
	}

	t.Run("truncated at every length", func(t *testing.T) {
		// Every prefix of a valid chunk must error rather than panic. This is
		// the shape a torn write takes on disk.
		for n := 0; n < len(buf.b); n++ {
			func() {
				defer func() {
					if r := recover(); r != nil {
						t.Fatalf("prefix of length %d panicked: %v", n, r)
					}
				}()
				if _, err := DecodeChunk(buf.b[:n]); err == nil {
					t.Fatalf("prefix of length %d decoded without error", n)
				}
			}()
		}
	})

	t.Run("bit flips", func(t *testing.T) {
		// CRC catches these on the real read path; this asserts the decoder
		// itself survives them, since a query must not die on one bad shard.
		rng := rand.New(rand.NewSource(1))
		for i := 0; i < 500; i++ {
			c := make([]byte, len(buf.b))
			copy(c, buf.b)
			c[rng.Intn(len(c))] ^= 1 << uint(rng.Intn(8))
			func() {
				defer func() {
					if r := recover(); r != nil {
						t.Fatalf("bit flip panicked: %v", r)
					}
				}()
				DecodeChunk(c) // error or success both fine; a panic is not
			}()
		}
	})
}

func TestBadMagicRejected(t *testing.T) {
	if _, err := DecodeChunk([]byte("NOPE\x01\x00")); err == nil {
		t.Fatal("decoded a chunk with bad magic")
	}
}