stream_test.go

package tok_test

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/GrayCodeAI/tok"
)

func TestStreamCompressor_BasicAppendAndSnapshot(t *testing.T) {
	sc := tok.NewStreamCompressor(0) // default threshold
	defer sc.Close()

	sc.Append("Hello world")
	sc.Append("This is a test")

	snapshot, _ := sc.Snapshot()
	if snapshot == "" {
		t.Fatal("expected non-empty snapshot")
	}

	raw := sc.Raw()
	if !strings.Contains(raw, "Hello world") {
		t.Error("raw should contain appended content")
	}
	if !strings.Contains(raw, "This is a test") {
		t.Error("raw should contain all appended content")
	}
}

func TestStreamCompressor_EmptyAppend(t *testing.T) {
	sc := tok.NewStreamCompressor(100)
	defer sc.Close()

	sc.Append("")
	raw := sc.Raw()
	if raw != "" {
		t.Errorf("expected empty raw after appending empty string, got %q", raw)
	}
}

func TestStreamCompressor_ThresholdTriggersCompression(t *testing.T) {
	// Use a very low threshold so compression triggers quickly.
	sc := tok.NewStreamCompressor(10, tok.Aggressive)
	defer sc.Close()

	// Append enough content to exceed the threshold.
	bigContent := strings.Repeat("The quick brown fox jumped over the lazy dog. ", 50)
	sc.Append(bigContent)

	// Wait for background compression to complete.
	time.Sleep(500 * time.Millisecond)

	snapshot, stats := sc.Snapshot()
	if snapshot == "" {
		t.Fatal("expected non-empty compressed snapshot after threshold exceeded")
	}
	if stats.OriginalTokens == 0 {
		t.Error("expected non-zero original tokens in stats after compression")
	}
}

func TestStreamCompressor_SnapshotBeforeCompression(t *testing.T) {
	// High threshold so compression never triggers.
	sc := tok.NewStreamCompressor(999999)
	defer sc.Close()

	sc.Append("some content")

	snapshot, stats := sc.Snapshot()
	if snapshot == "" {
		t.Fatal("expected non-empty snapshot (should return raw)")
	}
	if stats.OriginalTokens != 0 {
		t.Error("expected zero stats when compression hasn't run")
	}
}

func TestStreamCompressor_TokenCount(t *testing.T) {
	sc := tok.NewStreamCompressor(999999)
	defer sc.Close()

	sc.Append("hello world this is a token counting test")
	count := sc.TokenCount()
	if count == 0 {
		t.Error("expected non-zero token count")
	}
}

func TestStreamCompressor_ConcurrentAccess(t *testing.T) {
	sc := tok.NewStreamCompressor(50)
	defer sc.Close()

	var wg sync.WaitGroup
	// Multiple goroutines appending concurrently.
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			sc.Append("concurrent content that is long enough to matter for tokens. ")
		}()
	}

	// Multiple goroutines reading snapshots concurrently.
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			snapshot, _ := sc.Snapshot()
			_ = snapshot
		}()
	}

	// Multiple goroutines reading raw concurrently.
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_ = sc.Raw()
		}()
	}

	// Multiple goroutines reading token count concurrently.
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_ = sc.TokenCount()
		}()
	}

	wg.Wait()
}

func TestStreamCompressor_NoDoubleCompression(t *testing.T) {
	// Verify rapid Append calls don't spawn multiple goroutines.
	// We use a low threshold to ensure compression would trigger.
	sc := tok.NewStreamCompressor(5, tok.Minimal)
	defer sc.Close()

	// Rapidly append many times; only one compression should be in progress at a time.
	for i := 0; i < 50; i++ {
		sc.Append("word word word word word word word word word word ")
	}

	// Allow time for background compression.
	time.Sleep(500 * time.Millisecond)

	snapshot, _ := sc.Snapshot()
	if snapshot == "" {
		t.Fatal("expected non-empty snapshot after many appends")
	}
}

func TestStreamCompressor_CloseWaitsForCompression(t *testing.T) {
	sc := tok.NewStreamCompressor(5, tok.Aggressive)

	// Append enough to trigger compression.
	bigContent := strings.Repeat("token token token token token ", 100)
	sc.Append(bigContent)

	// Close should wait for any in-progress compression and not panic.
	sc.Close()

	// After close, Snapshot should still return valid data.
	snapshot, _ := sc.Snapshot()
	if snapshot == "" {
		t.Fatal("expected non-empty snapshot after close")
	}
}

func TestStreamCompressor_DefaultThreshold(t *testing.T) {
	// Passing 0 should default to 2000.
	sc := tok.NewStreamCompressor(0)
	defer sc.Close()

	// Small content should not trigger compression.
	sc.Append("small text")
	time.Sleep(100 * time.Millisecond)

	_, stats := sc.Snapshot()
	if stats.OriginalTokens != 0 {
		t.Error("expected no compression for content well below default threshold")
	}
}