218 lines
7.6 KiB
Go
218 lines
7.6 KiB
Go
package pipeline
|
||
|
||
import (
|
||
"reflect"
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
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 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)
|
||
}
|
||
}
|
||
|
||
// 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
|
||
}
|