305 lines
11 KiB
Go
305 lines
11 KiB
Go
package pipeline
|
||
|
||
import (
|
||
"reflect"
|
||
"strings"
|
||
"testing"
|
||
"unicode"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
func TestSplitChunksSingleParagraph(t *testing.T) {
|
||
got := SplitChunks([]string{"静かな図書館の朝。"})
|
||
want := []Chunk{{Chapter: 1, ChunkIdx: 0, Text: "静かな図書館の朝。"}}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("single paragraph = %+v, want %+v", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksChapters(t *testing.T) {
|
||
got := SplitChunks([]string{"Глава один.", "Глава два.", "Глава три."})
|
||
want := []Chunk{
|
||
{Chapter: 1, ChunkIdx: 0, Text: "Глава один."},
|
||
{Chapter: 2, ChunkIdx: 0, Text: "Глава два."},
|
||
{Chapter: 3, ChunkIdx: 0, Text: "Глава три."},
|
||
}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("chapters = %+v, want %+v", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksPacksToTarget(t *testing.T) {
|
||
// Two paragraphs that together exceed the TOKEN target (each ~1000 tokens for
|
||
// cyrillic: 3000 chars / 3) must split into two chunks, each kept whole.
|
||
p1 := strings.Repeat("а", 3000)
|
||
p2 := strings.Repeat("б", 3000)
|
||
got := SplitChunks([]string{p1 + "\n\n" + p2})
|
||
want := []Chunk{
|
||
{Chapter: 1, ChunkIdx: 0, Text: p1},
|
||
{Chapter: 1, ChunkIdx: 1, Text: p2},
|
||
}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("packing = %d chunks, want 2 whole paragraphs; got %+v", len(got), summarize(got))
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksPacksSmallParagraphsTogether(t *testing.T) {
|
||
got := SplitChunks([]string{"Абзац один.\n\nАбзац два.\n\nАбзац три."})
|
||
if len(got) != 1 {
|
||
t.Fatalf("small paragraphs must pack into one chunk, got %d: %+v", len(got), summarize(got))
|
||
}
|
||
if got[0].Text != "Абзац один.\n\nАбзац два.\n\nАбзац три." {
|
||
t.Fatalf("packed text = %q", got[0].Text)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksOversizeParagraphSplitsAtSentences(t *testing.T) {
|
||
// One paragraph far over the token target, built of many whole sentences, must
|
||
// split into >1 chunk AND never cut a sentence: every original sentence must
|
||
// survive intact and in order across the chunks.
|
||
sentence := strings.Repeat("слово ", 40) + "конец." // ~80 tokens each
|
||
var sb strings.Builder
|
||
for i := 0; i < 40; i++ { // ~3200 tokens total ≫ 1500
|
||
sb.WriteString(sentence)
|
||
sb.WriteByte(' ')
|
||
}
|
||
para := strings.TrimSpace(sb.String())
|
||
got := SplitChunks([]string{para})
|
||
if len(got) < 2 {
|
||
t.Fatalf("oversize paragraph must split into >1 chunk, got %d", len(got))
|
||
}
|
||
wantSentences := trimAll(splitSourceSentences(para))
|
||
var gotSentences []string
|
||
for _, c := range got {
|
||
gotSentences = append(gotSentences, trimAll(splitSourceSentences(c.Text))...)
|
||
}
|
||
if !reflect.DeepEqual(gotSentences, wantSentences) {
|
||
t.Fatalf("sentence integrity broken:\n got %d sentences\nwant %d sentences", len(gotSentences), len(wantSentences))
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksShortLinesPackByTrueTokenBudget(t *testing.T) {
|
||
// Many short dialogue paragraphs must pack by TRUE token content, not by summing
|
||
// the per-unit 16-token floor. 150 lines of 「はい。」 are ~450 true tokens (cjk=300,
|
||
// other=450 → 300+150), so they belong in ONE chunk. Summing the floor (150×16)
|
||
// would wrongly flush into ~2 sub-500-char chunks the coverage gate then skips.
|
||
var paras []string
|
||
for i := 0; i < 150; i++ {
|
||
paras = append(paras, "「はい。」")
|
||
}
|
||
got := SplitChunks([]string{strings.Join(paras, "\n\n")})
|
||
if len(got) != 1 {
|
||
t.Fatalf("short lines must pack by true token budget into 1 chunk, got %d chunks", len(got))
|
||
}
|
||
if EstimateTokens(got[0].Text) > targetChunkTokens {
|
||
t.Fatalf("packed chunk overshoots the budget: %d tokens", EstimateTokens(got[0].Text))
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksOversizeSentenceIsOwnChunk(t *testing.T) {
|
||
// A single sentence (no internal terminator) larger than the budget can only be
|
||
// its own chunk — never split.
|
||
big := strings.Repeat("ы", 6000) // ~2000 tokens, one un-terminated run
|
||
got := SplitChunks([]string{big})
|
||
if len(got) != 1 || got[0].Text != big {
|
||
t.Fatalf("an oversize sentence must be one un-split chunk, got %+v", summarize(got))
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksDropsEmpties(t *testing.T) {
|
||
// A blank/whitespace-only chapter vanishes and does NOT consume a chapter number.
|
||
got := SplitChunks([]string{"A", " \n\n ", "B"})
|
||
want := []Chunk{
|
||
{Chapter: 1, ChunkIdx: 0, Text: "A"},
|
||
{Chapter: 2, ChunkIdx: 0, Text: "B"},
|
||
}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("empties = %+v, want %+v", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksDeterministic(t *testing.T) {
|
||
chapters := []string{
|
||
"Абзац.\n\nЕщё абзац.",
|
||
strings.Repeat("слово ", 400) + "конец.",
|
||
}
|
||
a, b := SplitChunks(chapters), SplitChunks(chapters)
|
||
if !reflect.DeepEqual(a, b) {
|
||
t.Fatal("SplitChunks must be deterministic")
|
||
}
|
||
}
|
||
|
||
func TestSplitChunksNeverSplitsAbbreviation(t *testing.T) {
|
||
// Two long en paragraphs each ending mid-flow with "Mr. Smith" must not chunk
|
||
// between "Mr." and "Smith" — the abbreviation guard keeps the sentence whole.
|
||
p := strings.Repeat("The road went on and on. ", 100) + "It was Mr. Smith who arrived."
|
||
got := SplitChunks([]string{p})
|
||
for _, c := range got {
|
||
if strings.HasSuffix(strings.TrimSpace(c.Text), "Mr.") {
|
||
t.Fatalf("chunk ended at an abbreviation 'Mr.' — a sentence was cut: %q", c.Text)
|
||
}
|
||
}
|
||
}
|
||
|
||
// --- source sentence splitter (zh/ja/en) ---------------------------------------
|
||
|
||
func TestSplitSourceSentencesTilesLosslessly(t *testing.T) {
|
||
// The segments must TILE the input: concatenation reproduces it byte-for-byte.
|
||
inputs := []string{
|
||
"静かな朝。彼は歩いた。「止まれ！」と叫んだ。",
|
||
"他说。“不要！”然后离开了。",
|
||
"Hello world. This is a test! Really? Yes.",
|
||
"Mr. Smith met Dr. J. R. R. Brown. They talked.",
|
||
"末尾に句点なし",
|
||
"",
|
||
}
|
||
for _, in := range inputs {
|
||
segs := splitSourceSentences(in)
|
||
if got := strings.Join(segs, ""); got != in {
|
||
t.Fatalf("tiling broke for %q: rejoined to %q", in, got)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesJapanese(t *testing.T) {
|
||
// CJK terminators split zero-width; a closing quote is absorbed into its sentence.
|
||
got := trimAll(splitSourceSentences("静かな朝。「止まれ！」と彼は言った。"))
|
||
want := []string{"静かな朝。", "「止まれ！」と彼は言った。"}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("ja split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesEnglish(t *testing.T) {
|
||
got := trimAll(splitSourceSentences(`He said "Stop!" She ran. Done.`))
|
||
want := []string{`He said "Stop!"`, "She ran.", "Done."}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("en split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesEllipsisNotOverSplitInCJK(t *testing.T) {
|
||
// "……" mid-CJK-sentence must not split; the trailing 。 still ends it.
|
||
got := trimAll(splitSourceSentences("そう……だった。次へ。"))
|
||
want := []string{"そう……だった。", "次へ。"}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("ellipsis split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesASCIIRespectsQuoteDepth(t *testing.T) {
|
||
// ASCII terminators inside a tracked quote must NOT split (like the CJK branch):
|
||
// 「Stop. Now.」 is one dialogue unit, not two sentences (external-review #4).
|
||
got := trimAll(splitSourceSentences("「Stop. Now.」次へ。"))
|
||
want := []string{"「Stop. Now.」次へ。"}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("ASCII-in-quote split = %q, want %q (must not over-split dialogue)", got, want)
|
||
}
|
||
// Sanity: outside quotes ASCII still splits.
|
||
if n := len(trimAll(splitSourceSentences("Stop. Now."))); n != 2 {
|
||
t.Fatalf("plain ASCII must still split into 2, got %d", n)
|
||
}
|
||
}
|
||
|
||
func TestSplitSourceSentencesAbbreviations(t *testing.T) {
|
||
// Abbreviations and single-letter initials do not end a sentence.
|
||
got := trimAll(splitSourceSentences("Dr. J. R. R. Tolkien wrote it. The end."))
|
||
want := []string{"Dr. J. R. R. Tolkien wrote it.", "The end."}
|
||
if !reflect.DeepEqual(got, want) {
|
||
t.Fatalf("abbrev split = %q, want %q", got, want)
|
||
}
|
||
}
|
||
|
||
// --- property/fuzz (external-review #7: would have caught #3) ----------------
|
||
|
||
// FuzzSplitSourceSentencesTiles asserts the segmenter is lossless: its segments
|
||
// always concatenate back to the input's rune sequence, and byte-exactly when the
|
||
// input is valid UTF-8 (the production invariant after NormalizeSource). A path that
|
||
// dropped/reordered/replaced a rune — like the size-dependent []rune→U+FFFD bug #3 —
|
||
// fails this over random input.
|
||
func FuzzSplitSourceSentencesTiles(f *testing.F) {
|
||
for _, s := range []string{
|
||
"静かな朝。「止まれ！」と言った。", "Hello. World!", "Mr. Smith went. Home.",
|
||
"そう……だった。", "「Stop. Now.」", "abc\xe9def", "",
|
||
} {
|
||
f.Add(s)
|
||
}
|
||
f.Fuzz(func(t *testing.T, s string) {
|
||
segs := splitSourceSentences(s)
|
||
joined := strings.Join(segs, "")
|
||
if joined != string([]rune(s)) {
|
||
t.Fatalf("tiling lost/reordered runes: input=%q rejoined=%q", s, joined)
|
||
}
|
||
if utf8.ValidString(s) && joined != s {
|
||
t.Fatalf("valid-UTF-8 tiling not byte-exact: %q → %q", s, joined)
|
||
}
|
||
})
|
||
}
|
||
|
||
// FuzzNormalizeSourceValidUTF8 guards fix #3 directly: NormalizeSource always yields
|
||
// valid UTF-8, so the chunker never sees a stray byte the two paths would treat
|
||
// differently.
|
||
func FuzzNormalizeSourceValidUTF8(f *testing.F) {
|
||
for _, s := range []string{"abc\xe9def", "正常なテキスト", "\xff\xfe", ""} {
|
||
f.Add(s)
|
||
}
|
||
f.Fuzz(func(t *testing.T, s string) {
|
||
if out := NormalizeSource(s); !utf8.ValidString(out) {
|
||
t.Fatalf("NormalizeSource must return valid UTF-8, invalid for %q", s)
|
||
}
|
||
})
|
||
}
|
||
|
||
// FuzzSplitChunksPreservesContent asserts chunking loses no non-space content: the
|
||
// count of non-space runes across all chunks equals that of the normalized source.
|
||
func FuzzSplitChunksPreservesContent(f *testing.F) {
|
||
for _, s := range []string{
|
||
"静かな朝。\n\n「止まれ！」と彼は言った。次へ。", strings.Repeat("слово ", 500), "abc\xe9def", "",
|
||
} {
|
||
f.Add(s)
|
||
}
|
||
f.Fuzz(func(t *testing.T, s string) {
|
||
norm := NormalizeSource(s)
|
||
want := countNonSpace(norm)
|
||
got := 0
|
||
for _, c := range SplitChunks([]string{norm}) {
|
||
got += countNonSpace(c.Text)
|
||
}
|
||
if got != want {
|
||
t.Fatalf("chunking lost non-space content: got %d want %d (input %q)", got, want, s)
|
||
}
|
||
})
|
||
}
|
||
|
||
func countNonSpace(s string) int {
|
||
n := 0
|
||
for _, r := range s {
|
||
if !unicode.IsSpace(r) {
|
||
n++
|
||
}
|
||
}
|
||
return n
|
||
}
|
||
|
||
// trimAll TrimSpaces each segment and drops empties (test convenience).
|
||
func trimAll(segs []string) []string {
|
||
var out []string
|
||
for _, s := range segs {
|
||
if t := strings.TrimSpace(s); t != "" {
|
||
out = append(out, t)
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func summarize(cs []Chunk) []string {
|
||
out := make([]string, len(cs))
|
||
for i, c := range cs {
|
||
txt := c.Text
|
||
if len(txt) > 12 {
|
||
txt = txt[:12] + "…"
|
||
}
|
||
out[i] = strings.TrimSpace(txt)
|
||
}
|
||
return out
|
||
}
|