cost-model/core/pkg/reader/jsonlines_test.go

package reader

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand/v2"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"github.com/google/uuid"
)

// jsonLinesOf marshals each of vals to its own line, returning a JSONL reader.
func jsonLinesOf[T any](t *testing.T, vals []T, delimiter string) io.Reader {
	t.Helper()
	var sb strings.Builder
	for _, v := range vals {
		b, err := json.Marshal(v)
		if err != nil {
			t.Fatalf("marshaling test data: %v", err)
		}
		sb.Write(b)
		if delimiter != "" {
			sb.WriteString(delimiter)
		}
	}
	return strings.NewReader(sb.String())
}

// recordingCloser wraps a reader and records whether Close was called, returning
// a configurable error from Close. Used to exercise the io.Closer branch.
type recordingCloser struct {
	io.Reader
	closed   bool
	closeErr error
}

func (rc *recordingCloser) Close() error {
	rc.closed = true
	return rc.closeErr
}

// TestJSONLinesReader_FoldedTerminalSignal proves the JSONL reader honors the
// same folded-io.EOF contract as the array reader (folds when the final batch is
// short; defers to (0, io.EOF) on an exact fill).
func TestJSONLinesReader_FoldedTerminalSignal(t *testing.T) {
	tests := []struct {
		name      string
		items     int
		bufSize   int
		wantReads []struct {
			n   int
			eof bool
		}
	}{
		{"buffer larger than items", 3, 10, []struct {
			n   int
			eof bool
		}{{3, true}}},
		{"buffer exactly fits items", 5, 5, []struct {
			n   int
			eof bool
		}{{5, false}, {0, true}}},
		{"buffer unevenly divides items", 7, 3, []struct {
			n   int
			eof bool
		}{{3, false}, {3, false}, {1, true}}},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			r := NewJSONLinesReader[int](jsonLinesOf(t, seq(tc.items), "\n"))
			dst := make([]int, tc.bufSize)
			for i, want := range tc.wantReads {
				n, err := r.Read(context.Background(), dst)
				if n != want.n {
					t.Errorf("read %d: got n=%d, want %d", i, n, want.n)
				}
				if gotEOF := errors.Is(err, io.EOF); gotEOF != want.eof {
					t.Errorf("read %d: got eof=%v (err=%v), want eof=%v", i, gotEOF, err, want.eof)
				}
				if !want.eof && err != nil {
					t.Errorf("read %d: unexpected error: %v", i, err)
				}
			}
		})
	}
}

// TestJSONLinesReader_Whitespace verifies newline separation, blank lines,
// trailing whitespace, CRLF, and empty input are all handled.
func TestJSONLinesReader_Whitespace(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  []int
	}{
		{"one per line", "1\n2\n3\n", []int{1, 2, 3}},
		{"no trailing newline", "1\n2\n3", []int{1, 2, 3}},
		{"space separated", "1 2 3", []int{1, 2, 3}},
		{"crlf", "1\r\n2\r\n3\r\n", []int{1, 2, 3}},
		{"blank lines", "1\n\n\n2\n", []int{1, 2}},
		{"trailing whitespace", "1\n2\n   \n", []int{1, 2}},
		{"empty", "", nil},
		{"all whitespace", "   \n  ", nil},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			r := NewJSONLinesReader[int](strings.NewReader(tc.input))
			dst := make([]int, 2)
			var got []int
			for {
				n, err := r.Read(context.Background(), dst)
				got = append(got, dst[:n]...)
				if errors.Is(err, io.EOF) {
					break
				}
				if err != nil {
					t.Fatalf("unexpected error: %v", err)
				}
			}
			if len(got) != len(tc.want) {
				t.Fatalf("got %v, want %v", got, tc.want)
			}
			for i := range got {
				if got[i] != tc.want[i] {
					t.Errorf("item %d: got %d, want %d", i, got[i], tc.want[i])
				}
			}
		})
	}
}

// TestJSONLinesReader_TruncatedValue verifies a truncated trailing value is a
// terminal, sticky, non-EOF error returned after the valid prefix — unlike the
// array reader, whose lookahead swallows truncation.
func TestJSONLinesReader_TruncatedValue(t *testing.T) {
	r := NewJSONLinesReader[int](strings.NewReader("1\n2\n{\"id\":"))
	dst := make([]int, 2)

	if n, err := r.Read(context.Background(), dst); n != 2 || err != nil {
		t.Fatalf("first read: got (%d, %v), want (2, nil)", n, err)
	}

	n, err := r.Read(context.Background(), dst)
	if n != 0 || err == nil || errors.Is(err, io.EOF) {
		t.Fatalf("second read: got (%d, %v), want (0, non-EOF error)", n, err)
	}
	// Sticky.
	if _, err2 := r.Read(context.Background(), dst); err2 != err {
		t.Errorf("third read: err=%v, want sticky %v", err2, err)
	}
}

// TestJSONLinesReader_MalformedValue verifies a wrong-type value mid-stream is a
// terminal, sticky, non-EOF error.
func TestJSONLinesReader_MalformedValue(t *testing.T) {
	r := NewJSONLinesReader[int](strings.NewReader("1\n{}\n3"))
	dst := make([]int, 4)

	// The batch stops at the malformed second value, returning the prefix.
	n, err := r.Read(context.Background(), dst)
	if n != 1 || err == nil || errors.Is(err, io.EOF) {
		t.Fatalf("read: got (%d, %v), want (1, non-EOF error)", n, err)
	}
	if dst[0] != 1 {
		t.Errorf("prefix item: got %d, want 1", dst[0])
	}
	if _, err2 := r.Read(context.Background(), dst); err2 != err {
		t.Errorf("next read: err=%v, want sticky %v", err2, err)
	}
}

// TestJSONLinesReader_StructValues exercises the realistic case: one JSON object
// per line decoded into pointer elements across multiple batches.
func TestJSONLinesReader_StructValues(t *testing.T) {
	type item struct {
		ID   int    `json:"id"`
		Name string `json:"name"`
	}
	src := []*item{{1, "a"}, {2, "b"}, {3, "c"}, {4, "d"}, {5, "e"}}

	r := NewJSONLinesReader[*item](jsonLinesOf(t, src, "\n"))
	dst := make([]*item, 2)

	var got []*item
	for {
		n, err := r.Read(context.Background(), dst)
		got = append(got, dst[:n]...)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
	}

	if len(got) != len(src) {
		t.Fatalf("got %d items, want %d", len(got), len(src))
	}
	for i, p := range got {
		if p == nil || p.ID != src[i].ID || p.Name != src[i].Name {
			t.Errorf("item %d: got %v, want %+v", i, p, *src[i])
		}
	}
}

// TestJSONLinesReader_StructValues_NoNewline tests a file without newlines
// between streamed JSON objects
func TestJSONLinesReader_StructValues_NoNewline(t *testing.T) {
	type item struct {
		ID   int    `json:"id"`
		Name string `json:"name"`
	}
	src := []*item{{1, "a"}, {2, "b"}, {3, "c"}, {4, "d"}, {5, "e"}}

	r := NewJSONLinesReader[*item](jsonLinesOf(t, src, ""))
	dst := make([]*item, 2)

	var got []*item
	for {
		n, err := r.Read(context.Background(), dst)
		got = append(got, dst[:n]...)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
	}

	if len(got) != len(src) {
		t.Fatalf("got %d items, want %d", len(got), len(src))
	}
	for i, p := range got {
		if p == nil || p.ID != src[i].ID || p.Name != src[i].Name {
			t.Errorf("item %d: got %v, want %+v", i, p, *src[i])
		}
	}
}

// TestJSONLinesReader_Close verifies the io.Closer source is closed.
func TestJSONLinesReader_Close(t *testing.T) {
	rc := &recordingCloser{Reader: strings.NewReader("1\n2\n")}
	r := NewJSONLinesReader[int](rc)
	if err := r.Close(); err != nil {
		t.Errorf("Close: got err=%v, want nil", err)
	}
	if !rc.closed {
		t.Error("underlying source was not closed")
	}
}

// Simple struct with various data types for benchmarking JSONLinesReader
type benchRecord struct {
	ID     string            `json:"id"`
	Name   string            `json:"name"`
	Value  float64           `json:"value"`
	Labels map[string]string `json:"labels"`
}

// writeBenchJSONL writes n JSON-Lines records to path (one value per line).
func writeBenchJSONL(tb testing.TB, path string, n int) {
	tb.Helper()

	f, err := os.Create(path)
	if err != nil {
		tb.Fatal(err)
	}
	defer f.Close()

	w := bufio.NewWriter(f)
	enc := json.NewEncoder(w) // Encode appends a newline after each value.
	for i := 0; i < n; i++ {
		rec := benchRecord{
			ID:     uuid.NewString(),
			Name:   fmt.Sprintf("name-%d", i),
			Value:  rand.Float64(),
			Labels: map[string]string{"foo": "bar", "baz": "bat"},
		}
		if err := enc.Encode(&rec); err != nil {
			tb.Fatal(err)
		}
	}
	if err := w.Flush(); err != nil {
		tb.Fatal(err)
	}
}

// streamAll reads path in batches of batchSize, discarding every batch (dst is
// reused, nothing retained), and returns the number of records seen.
func streamAll(tb testing.TB, path string, batchSize int) int {
	tb.Helper()

	f, err := os.Open(path)
	if err != nil {
		tb.Fatal(err)
	}

	r := NewJSONLinesReader[benchRecord](f)
	dst := make([]benchRecord, batchSize)
	total := 0
	for {
		n, err := r.Read(context.Background(), dst)
		total += n
		// Discard: the next Read overwrites dst; no record outlives its batch.
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			tb.Fatal(err)
		}
	}
	if err := r.Close(); err != nil {
		tb.Fatal(err)
	}
	return total
}

const (
	benchRecords   = 100_000
	benchBatchSize = 1_000
)

// BenchmarkJSONLinesReader measures wall time and allocation churn to stream a
// 50k-line JSONL file 1000 records at a time. With -benchmem, B/op is the TOTAL
// bytes allocated per pass (churn, scales with record count) — not resident
// memory. See BenchmarkJSONLinesReaderResident for the live-heap bound.
func BenchmarkJSONLinesReader(b *testing.B) {
	path := filepath.Join(b.TempDir(), "bench.jsonl") // disposable; auto-removed
	writeBenchJSONL(b, path, benchRecords)

	b.ReportAllocs()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		if got := streamAll(b, path, benchBatchSize); got != benchRecords {
			b.Fatalf("streamed %d records, want %d", got, benchRecords)
		}
	}
}

// BenchmarkJSONLinesReaderResident measures PEAK live heap (HeapInuse) while
// streaming, which is what the streaming design is meant to bound: it should
// stay ~flat at one batch's worth of records regardless of file size. Note the
// per-batch runtime.ReadMemStats sampling perturbs timing, so read ns/op from
// BenchmarkJSONLinesReader, not this one.
func BenchmarkJSONLinesReaderResident(b *testing.B) {
	path := filepath.Join(b.TempDir(), "bench.jsonl")
	writeBenchJSONL(b, path, benchRecords)

	b.ResetTimer()

	var peakHeapInuse uint64
	for i := 0; i < b.N; i++ {
		f, err := os.Open(path)
		if err != nil {
			b.Fatal(err)
		}
		r := NewJSONLinesReader[benchRecord](f)
		dst := make([]benchRecord, benchBatchSize)
		for {
			n, err := r.Read(context.Background(), dst)
			_ = n

			var ms runtime.MemStats
			runtime.ReadMemStats(&ms)
			if ms.HeapInuse > peakHeapInuse {
				peakHeapInuse = ms.HeapInuse
			}

			if errors.Is(err, io.EOF) {
				break
			}
			if err != nil {
				b.Fatal(err)
			}
		}
		r.Close()
	}

	b.ReportMetric(float64(peakHeapInuse)/1024, "peakHeapKB")
}