markovbot/markov/chain_test.go

package markov

import (
	"bytes"
	"errors"
	"math"
	"math/rand"
	"path/filepath"
	"strings"
	"testing"
)

func mustNew(t *testing.T, order int) *Chain {
	t.Helper()
	c, err := New(order)
	if err != nil {
		t.Fatalf("New(%d): %v", order, err)
	}
	return c
}

func TestNewRejectsBadOrder(t *testing.T) {
	for _, order := range []int{0, -1, -100} {
		if _, err := New(order); err == nil {
			t.Errorf("New(%d) should fail", order)
		}
	}
}

func TestGenerateUntrained(t *testing.T) {
	c := mustNew(t, 2)
	_, err := c.Generate(rand.New(rand.NewSource(1)), 10)
	if !errors.Is(err, ErrUntrained) {
		t.Errorf("err = %v, want ErrUntrained", err)
	}
}

// A chain trained on exactly one sentence has no choice at any step, so it
// must reproduce that sentence and then stop. This pins training,
// generation and termination in one assertion.
func TestSingleSampleIsReproduced(t *testing.T) {
	for order := 1; order <= 3; order++ {
		c := mustNew(t, order)
		const want = "the quick brown fox jumps over the lazy dog"
		c.Train(want)
		got, err := c.Generate(rand.New(rand.NewSource(1)), 100)
		if err != nil {
			t.Fatalf("order %d: %v", order, err)
		}
		if order > 1 && got != want {
			// At order 1 "the" is ambiguous, so only higher orders are
			// guaranteed to reproduce the input exactly.
			t.Errorf("order %d: got %q, want %q", order, got, want)
		}
		if order == 1 && got == "" {
			t.Errorf("order 1 produced nothing")
		}
	}
}

func TestGenerationTerminates(t *testing.T) {
	c := mustNew(t, 2)
	c.Train("a b c")
	got, err := c.Generate(rand.New(rand.NewSource(1)), 1000)
	if err != nil {
		t.Fatal(err)
	}
	// It must stop at the end sentinel, well short of the 1000 cap.
	if n := len(strings.Fields(got)); n != 3 {
		t.Errorf("got %d words (%q), want 3 — generation did not stop at the sample end", n, got)
	}
}

func TestMaxWordsIsRespected(t *testing.T) {
	c := mustNew(t, 1)
	// A cycle: without a cap this would generate forever.
	c.Train("a b a b a b a b a b")
	for _, max := range []int{1, 3, 7} {
		got, err := c.Generate(rand.New(rand.NewSource(1)), max)
		if err != nil {
			t.Fatal(err)
		}
		if n := len(strings.Fields(got)); n > max {
			t.Errorf("maxWords=%d produced %d words: %q", max, n, got)
		}
	}
}

func TestGenerateRejectsBadArgs(t *testing.T) {
	c := mustNew(t, 2)
	c.Train("hello world")
	if _, err := c.Generate(nil, 10); err == nil {
		t.Error("nil rng should fail")
	}
	if _, err := c.Generate(rand.New(rand.NewSource(1)), 0); err == nil {
		t.Error("maxWords 0 should fail")
	}
}

// Same seed must give the same text. This is the property Go's randomised
// map iteration would silently break if sampling walked the map directly.
func TestDeterministicForSeed(t *testing.T) {
	c := mustNew(t, 2)
	c.TrainAll([]string{
		"the cat sat on the mat",
		"the dog sat on the log",
		"the cat ate the rat",
		"a bird flew over the house",
	})
	for seed := int64(0); seed < 25; seed++ {
		first, err := c.Generate(rand.New(rand.NewSource(seed)), 30)
		if err != nil {
			t.Fatal(err)
		}
		for range 5 {
			again, err := c.Generate(rand.New(rand.NewSource(seed)), 30)
			if err != nil {
				t.Fatal(err)
			}
			if again != first {
				t.Fatalf("seed %d not reproducible:\n %q\n %q", seed, first, again)
			}
		}
	}
}

// Successors must be drawn in proportion to how often they were observed.
// Training "a b" three times and "a c" once should yield roughly 75/25.
func TestSamplingFollowsObservedFrequency(t *testing.T) {
	c := mustNew(t, 1)
	for range 300 {
		c.Train("a b")
	}
	for range 100 {
		c.Train("a c")
	}

	rng := rand.New(rand.NewSource(42))
	counts := map[string]int{}
	const trials = 20000
	for range trials {
		got, err := c.GenerateFrom(rng, "a", 1)
		if err != nil {
			t.Fatal(err)
		}
		counts[got]++
	}

	gotRatio := float64(counts["b"]) / float64(trials)
	if math.Abs(gotRatio-0.75) > 0.03 {
		t.Errorf("P(b|a) = %.3f, want ~0.75 (counts: %v)", gotRatio, counts)
	}
}

func TestGenerateFromUnknownSeedFallsBack(t *testing.T) {
	c := mustNew(t, 2)
	c.Train("hello there world")
	got, err := c.GenerateFrom(rand.New(rand.NewSource(1)), "completely unseen tokens", 20)
	if err != nil {
		t.Fatalf("unknown seed should fall back, not fail: %v", err)
	}
	if got == "" {
		t.Error("fallback produced no output")
	}
}

func TestGenerateFromKnownSeedContinues(t *testing.T) {
	c := mustNew(t, 2)
	c.Train("the quick brown fox jumps")
	got, err := c.GenerateFrom(rand.New(rand.NewSource(1)), "quick brown", 10)
	if err != nil {
		t.Fatal(err)
	}
	if got != "fox jumps" {
		t.Errorf("got %q, want %q", got, "fox jumps")
	}
}

func TestTokenize(t *testing.T) {
	tests := []struct {
		in   string
		want []string
	}{
		{"a b c", []string{"a", "b", "c"}},
		{"  spaced   out  ", []string{"spaced", "out"}},
		{"tabs\tand\nnewlines", []string{"tabs", "and", "newlines"}},
		{"", nil},
		{"   ", nil},
	}
	for _, tt := range tests {
		got := Tokenize(tt.in)
		if len(got) != len(tt.want) {
			t.Errorf("Tokenize(%q) = %v, want %v", tt.in, got, tt.want)
			continue
		}
		for i := range got {
			if got[i] != tt.want[i] {
				t.Errorf("Tokenize(%q) = %v, want %v", tt.in, got, tt.want)
				break
			}
		}
	}
}

// A corpus containing the sentinel strings must not be able to inject
// chain boundaries.
func TestTokenizeStripsSentinels(t *testing.T) {
	got := Tokenize("real " + start + " tokens " + end + " only")
	want := []string{"real", "tokens", "only"}
	if len(got) != len(want) {
		t.Fatalf("got %v, want %v", got, want)
	}
	for i := range got {
		if got[i] != want[i] {
			t.Fatalf("got %v, want %v", got, want)
		}
	}
}

func TestEmptyTrainingIsIgnored(t *testing.T) {
	c := mustNew(t, 2)
	c.Train("")
	c.Train("    ")
	if c.Prefixes() != 0 {
		t.Errorf("empty input created %d prefixes", c.Prefixes())
	}
}

func TestSaveLoadRoundTrip(t *testing.T) {
	c := mustNew(t, 2)
	c.TrainAll([]string{
		"the cat sat on the mat",
		"the dog sat on the log",
		"a bird flew over the house",
	})

	var buf bytes.Buffer
	if err := c.Save(&buf); err != nil {
		t.Fatal(err)
	}
	loaded, err := Load(&buf)
	if err != nil {
		t.Fatal(err)
	}

	if loaded.Order() != c.Order() {
		t.Errorf("order = %d, want %d", loaded.Order(), c.Order())
	}
	if loaded.Prefixes() != c.Prefixes() {
		t.Errorf("prefixes = %d, want %d", loaded.Prefixes(), c.Prefixes())
	}
	// The reloaded chain must behave identically, not merely look similar.
	for seed := int64(0); seed < 20; seed++ {
		want, err := c.Generate(rand.New(rand.NewSource(seed)), 30)
		if err != nil {
			t.Fatal(err)
		}
		got, err := loaded.Generate(rand.New(rand.NewSource(seed)), 30)
		if err != nil {
			t.Fatal(err)
		}
		if got != want {
			t.Fatalf("seed %d: reloaded chain generated %q, want %q", seed, got, want)
		}
	}
}

func TestSaveFileAtomic(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "model.json")

	c := mustNew(t, 2)
	c.Train("hello world again")
	if err := c.SaveFile(path); err != nil {
		t.Fatal(err)
	}
	loaded, err := LoadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	if loaded.Prefixes() != c.Prefixes() {
		t.Errorf("prefixes = %d, want %d", loaded.Prefixes(), c.Prefixes())
	}

	// No temporary files may be left behind.
	matches, _ := filepath.Glob(filepath.Join(dir, "*.tmp*"))
	if len(matches) != 0 {
		t.Errorf("temp files left behind: %v", matches)
	}
}

func TestLoadRejectsBadInput(t *testing.T) {
	tests := []struct {
		name string
		in   string
	}{
		{"not json", "this is not json"},
		{"wrong version", `{"version":999,"order":2,"freq":{}}`},
		{"bad order", `{"version":1,"order":0,"freq":{}}`},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if _, err := Load(strings.NewReader(tt.in)); err == nil {
				t.Error("expected an error")
			}
		})
	}
}

// Totals are derived on load, so a model whose counts were edited by hand
// must still sample correctly rather than trusting a stale sum.
func TestLoadDerivesTotals(t *testing.T) {
	const raw = `{"version":1,"order":1,"freq":{"\u0000start":{"x":5},"x":{"\u0000end":1}}}`
	c, err := Load(strings.NewReader(raw))
	if err != nil {
		t.Fatal(err)
	}
	got, err := c.Generate(rand.New(rand.NewSource(1)), 5)
	if err != nil {
		t.Fatal(err)
	}
	if got != "x" {
		t.Errorf("got %q, want %q", got, "x")
	}
}

func TestConcurrentTrainAndGenerate(t *testing.T) {
	c := mustNew(t, 2)
	c.Train("seed the chain so generation has something to work with")

	done := make(chan struct{})
	for range 4 {
		go func() {
			defer func() { done <- struct{}{} }()
			for i := range 200 {
				c.Train("concurrent training sample number " + string(rune('a'+i%26)))
			}
		}()
	}
	for range 4 {
		go func() {
			defer func() { done <- struct{}{} }()
			rng := rand.New(rand.NewSource(7))
			for range 200 {
				if _, err := c.Generate(rng, 20); err != nil {
					t.Errorf("generate: %v", err)
					return
				}
			}
		}()
	}
	for range 8 {
		<-done
	}
}

func BenchmarkTrain(b *testing.B) {
	c, _ := New(2)
	line := "the quick brown fox jumps over the lazy dog every single time"
	b.ResetTimer()
	for b.Loop() {
		c.Train(line)
	}
}

func BenchmarkGenerate(b *testing.B) {
	c, _ := New(2)
	for i := range 1000 {
		c.Train("sample number " + string(rune('a'+i%26)) + " with some trailing words")
	}
	rng := rand.New(rand.NewSource(1))
	b.ResetTimer()
	for b.Loop() {
		if _, err := c.Generate(rng, 30); err != nil {
			b.Fatal(err)
		}
	}
}