mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-06 19:38:36 +08:00
943 lines
28 KiB
Go
943 lines
28 KiB
Go
package layout
|
||
|
||
import (
|
||
pdf "ragflow/internal/deepdoc/parser/pdf/type"
|
||
util "ragflow/internal/deepdoc/parser/pdf/util"
|
||
"strings"
|
||
"testing"
|
||
)
|
||
|
||
// ---- test helpers ----
|
||
|
||
func newTestTextBox(page int, x0, x1, top, bottom float64, text string) pdf.TextBox {
|
||
return pdf.TextBox{
|
||
PageNumber: page,
|
||
X0: x0,
|
||
X1: x1,
|
||
Top: top,
|
||
Bottom: bottom,
|
||
Text: text,
|
||
}
|
||
}
|
||
|
||
func TestAssignColumn(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 50, Text: "col0-left"},
|
||
{PageNumber: 0, X0: 55, Text: "col0-mid"},
|
||
{PageNumber: 0, X0: 400, Text: "col1"},
|
||
{PageNumber: 1, X0: 50, Text: "pg1-col0"},
|
||
}
|
||
result := AssignColumn(boxes, 3)
|
||
if len(result) != 4 {
|
||
t.Fatal("expected 4 boxes")
|
||
}
|
||
if result[0].ColID != result[1].ColID {
|
||
t.Error("boxes 0 and 1 (close x0) should be same column")
|
||
}
|
||
if result[0].ColID == result[2].ColID {
|
||
t.Error("boxes 0 and 2 (far apart) should be different columns")
|
||
}
|
||
}
|
||
|
||
func TestTextMerge(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, ColID: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, Text: "左半", LayoutType: "text", LayoutNo: "1"},
|
||
{PageNumber: 0, ColID: 0, X0: 252, X1: 550, Top: 100, Bottom: 112, Text: "右半", LayoutType: "text", LayoutNo: "1"},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
result := TextMerge(boxes, meanH, 3)
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box, got %d", len(result))
|
||
}
|
||
}
|
||
|
||
func TestTextMergeNoMerge_DiffLayout(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, ColID: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, Text: "text", LayoutType: "text", LayoutNo: "1"},
|
||
{PageNumber: 0, ColID: 0, X0: 252, X1: 550, Top: 100, Bottom: 112, Text: "table", LayoutType: "table", LayoutNo: "2"},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
result := TextMerge(boxes, meanH, 3)
|
||
if len(result) != 2 {
|
||
t.Error("table and text should not merge")
|
||
}
|
||
}
|
||
|
||
func TestFinalReadingOrderMerge(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 1, ColID: 1, Top: 50, Text: "pg1-col1"},
|
||
{PageNumber: 0, ColID: 0, Top: 100, Text: "pg0-col0"},
|
||
{PageNumber: 0, ColID: 0, Top: 50, Text: "pg0-col0-top"},
|
||
}
|
||
result := FinalReadingOrderMerge(boxes)
|
||
if result[0].Text != "pg0-col0-top" {
|
||
t.Errorf("first should be pg0-col0-top: %q", result[0].Text)
|
||
}
|
||
if result[2].Text != "pg1-col1" {
|
||
t.Errorf("last should be pg1-col1: %q", result[2].Text)
|
||
}
|
||
}
|
||
|
||
func TestContainsRune(t *testing.T) {
|
||
if !strings.ContainsRune("。？！", '。') {
|
||
t.Error("should find 。")
|
||
}
|
||
if strings.ContainsRune("abc", 'z') {
|
||
t.Error("should not find z")
|
||
}
|
||
}
|
||
|
||
func TestEndsWithOneOf(t *testing.T) {
|
||
if !endsWithOneOf("句子结束。", "。？！?") {
|
||
t.Error("should match 。")
|
||
}
|
||
if endsWithOneOf("no match", "。？！?") {
|
||
t.Error("should not match")
|
||
}
|
||
}
|
||
|
||
func TestDefaultConfig(t *testing.T) {
|
||
cfg := pdf.DefaultParserConfig()
|
||
if cfg.Zoom != 3 {
|
||
t.Error("default zoom should be 3")
|
||
}
|
||
if cfg.ToPage != -1 {
|
||
t.Error("default to_page should be -1")
|
||
}
|
||
}
|
||
|
||
func TestKmeans1D_Boundary(t *testing.T) {
|
||
t.Run("n equals k", func(t *testing.T) {
|
||
data := []float64{50.0, 400.0}
|
||
labels, centroids := util.KMeans1D(data, 2)
|
||
if len(centroids) != 2 {
|
||
t.Errorf("n=k=2: expected 2 centroids, got %d — BUG: n<=k early return gives only 1 centroid", len(centroids))
|
||
}
|
||
if len(centroids) == 2 && labels[0] == labels[1] {
|
||
t.Error("n=k=2: two distinct points should be in different clusters — BUG: all points assigned to same cluster")
|
||
}
|
||
})
|
||
|
||
t.Run("n less than k", func(t *testing.T) {
|
||
data := []float64{100.0, 200.0, 300.0}
|
||
labels, centroids := util.KMeans1D(data, 4)
|
||
if len(centroids) != 3 {
|
||
t.Errorf("n=3,k=4: expected 3 centroids (one per point), got %d — BUG: n<=k early return gives only 1 centroid", len(centroids))
|
||
}
|
||
// All 3 points should be in different clusters
|
||
seen := make(map[int]bool)
|
||
for _, l := range labels {
|
||
seen[l] = true
|
||
}
|
||
if len(seen) != 3 {
|
||
t.Errorf("n=3,k=4: expected 3 distinct clusters, got %d", len(seen))
|
||
}
|
||
})
|
||
|
||
t.Run("single point", func(t *testing.T) {
|
||
data := []float64{100.0}
|
||
labels, centroids := util.KMeans1D(data, 1)
|
||
if len(centroids) != 1 || centroids[0] != 100.0 {
|
||
t.Errorf("single point: unexpected centroids %v", centroids)
|
||
}
|
||
if labels[0] != 0 {
|
||
t.Errorf("single point: label should be 0, got %d", labels[0])
|
||
}
|
||
})
|
||
}
|
||
|
||
// ---- startsWithOneOf / NaiveVerticalMerge (Issue 1: 、 vs ,) ----
|
||
|
||
func TestStartsWithOneOf(t *testing.T) {
|
||
// Python's concatting start-of-line character set:
|
||
// "。；？！?"）),，、："
|
||
// Go's set matches Python exactly.
|
||
|
||
// Use the CORRECT Python set to document expected behavior.
|
||
pySet := "。；？！?\")),，、："
|
||
|
||
t.Run("ASCII comma", func(t *testing.T) {
|
||
// Python concatting set includes ASCII comma U+002C.
|
||
// Go's set has 、(U+3001) instead — BUG.
|
||
if !startsWithOneOf(", rest", pySet) {
|
||
t.Error("should match ASCII comma ','")
|
||
}
|
||
})
|
||
|
||
t.Run("Chinese dun comma", func(t *testing.T) {
|
||
if !startsWithOneOf("、rest", pySet) {
|
||
t.Error("should match Chinese dun comma '、'")
|
||
}
|
||
})
|
||
|
||
t.Run("fullwidth comma", func(t *testing.T) {
|
||
if !startsWithOneOf("，rest", pySet) {
|
||
t.Error("should match fullwidth comma '，'")
|
||
}
|
||
})
|
||
|
||
t.Run("fullwidth period", func(t *testing.T) {
|
||
if !startsWithOneOf("。rest", pySet) {
|
||
t.Error("should match fullwidth period '。'")
|
||
}
|
||
})
|
||
|
||
t.Run("Chinese text should not match", func(t *testing.T) {
|
||
if startsWithOneOf("你好世界", pySet) {
|
||
t.Error("should NOT match Chinese text")
|
||
}
|
||
})
|
||
|
||
t.Run("letter should not match", func(t *testing.T) {
|
||
if startsWithOneOf("A letter", pySet) {
|
||
t.Error("should NOT match letter")
|
||
}
|
||
})
|
||
|
||
t.Run("empty string", func(t *testing.T) {
|
||
if startsWithOneOf("", pySet) {
|
||
t.Error("should NOT match empty string")
|
||
}
|
||
})
|
||
|
||
// Verify the actual Go set matches Python.
|
||
t.Run("Go set matches ASCII comma", func(t *testing.T) {
|
||
goSet := "。；？！?\"）),，、："
|
||
if !startsWithOneOf(", rest", goSet) {
|
||
t.Error("Go's concatting set should match ASCII comma ','")
|
||
}
|
||
})
|
||
|
||
t.Run("Go set has 、once", func(t *testing.T) {
|
||
goSet := "。；？！?\"）),，、："
|
||
count := 0
|
||
for _, r := range goSet {
|
||
if r == '、' {
|
||
count++
|
||
}
|
||
}
|
||
if count != 1 {
|
||
t.Errorf("Go set should have 、once, got %d", count)
|
||
}
|
||
})
|
||
}
|
||
|
||
func TestNaiveVerticalMerge_CommaConcat(t *testing.T) {
|
||
// When next line starts with ASCII comma ',' (U+002C), Python merges
|
||
// vertically because ',' is in the concatting startsWithOneOf set.
|
||
// Go now matches Python exactly — should merge.
|
||
|
||
t.Run("next line starts with ASCII comma", func(t *testing.T) {
|
||
// ASCII comma ',' is in Python's concatting set, Go matches.
|
||
// When there's NO anti trigger, merge happens by default.
|
||
// The concatting feature is only needed when it must OVERRIDE an anti trigger.
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "这是第一句话",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: ", 这是第二句话",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box, got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("ASCII comma should override period anti (now fixed)", func(t *testing.T) {
|
||
// Python: previous line ends with "。" (anti), next line starts with ","
|
||
// (concatting). Concatting OVERRIDES anti → merge.
|
||
// Go now matches Python: ',' is in concatting set → merge.
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "前一句话结束。",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: ", 这是续行",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box (ASCII comma ',' should override period anti), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("next line starts with fullwidth comma — should merge", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "这是第一句话",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: "，这是第二句话",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box (next line starts with '，'), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("next line starts with period — should merge", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "前文内容",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: "。这是下一句",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box (next line starts with '。'), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("no concat, no anti, no detach — should merge (default)", func(t *testing.T) {
|
||
// Python's _naive_vertical_merge: merge is the DEFAULT.
|
||
// concatting overrides anti; anti + detach prevent merge.
|
||
// When none trigger, boxes merge.
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "这是第一句话",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: "这是第二句话",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
// Default merge — no anti, no detach, same layoutno, close gap.
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box (default merge when no anti/detach), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("detach — horizontally separated boxes", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 100, Top: 100, Bottom: 112,
|
||
Text: "左列文字",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 300, X1: 350, Top: 114, Bottom: 126,
|
||
Text: "。右列文字",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 50}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
// Even with '。' concat char, boxes are detached horizontally.
|
||
if len(result) != 2 {
|
||
t.Errorf("expected 2 boxes (horizontally detached), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("large vertical gap — anti", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "第一句话",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 200, Bottom: 212,
|
||
Text: "。第二句话",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
// Gap 200-112=88 > 12*1.5=18 — anti triggers.
|
||
if len(result) != 2 {
|
||
t.Errorf("expected 2 boxes (large vertical gap), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("english period anti when isEnglish", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "End of sentence.",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: "Next sentence",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, true)
|
||
// When isEnglish=true, endsWith ".!?" is anti — don't merge.
|
||
if len(result) != 2 {
|
||
t.Errorf("expected 2 boxes (english period anti), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("cross-page — should NOT merge", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "第一页最后一行",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 1, X0: 50, X1: 250, Top: 50, Bottom: 62,
|
||
Text: "。第二页第一行",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
meanH := map[int]float64{0: 12, 1: 12}
|
||
meanW := map[int]float64{0: 200, 1: 200}
|
||
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
// Different pages — NaiveVerticalMerge groups by page.
|
||
if len(result) != 2 {
|
||
t.Errorf("expected 2 boxes (different pages), got %d", len(result))
|
||
}
|
||
})
|
||
|
||
t.Run("empty boxes", func(t *testing.T) {
|
||
result := NaiveVerticalMerge(nil, nil, nil, false)
|
||
if len(result) != 0 {
|
||
t.Error("expected empty result for nil input")
|
||
}
|
||
result = NaiveVerticalMerge([]pdf.TextBox{}, nil, nil, false)
|
||
if len(result) != 0 {
|
||
t.Error("expected empty result for empty input")
|
||
}
|
||
})
|
||
|
||
t.Run("single box", func(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112, Text: "only", LayoutNo: "1"},
|
||
}
|
||
result := NaiveVerticalMerge(boxes, nil, nil, false)
|
||
if len(result) != 1 {
|
||
t.Error("single box should be returned as-is")
|
||
}
|
||
})
|
||
}
|
||
|
||
// is applied.
|
||
func TestNaiveVerticalMerge_BottomShrink(t *testing.T) {
|
||
// Three boxes on the same page, sorted by Top.
|
||
// A + B merge first → tall box with Bottom=300.
|
||
// C overlaps vertically (Top=290 < prev.Bottom=300) but is short (Bottom=295).
|
||
// Current code: prev.Bottom = 295 (shrinks from 300).
|
||
// Correct: prev.Bottom = max(300, 295) = 300.
|
||
boxes := []pdf.TextBox{
|
||
{X0: 50, X1: 500, Top: 100, Bottom: 150, Text: "line one", PageNumber: 0},
|
||
{X0: 50, X1: 500, Top: 160, Bottom: 300, Text: "tall paragraph that spans many lines", PageNumber: 0},
|
||
{X0: 50, X1: 500, Top: 290, Bottom: 295, Text: "short overlap", PageNumber: 0},
|
||
}
|
||
mh := map[int]float64{0: 50} // threshold = 50 * 1.5 = 75
|
||
mw := map[int]float64{0: 5}
|
||
|
||
result := NaiveVerticalMerge(boxes, mh, mw, false)
|
||
|
||
if len(result) != 1 {
|
||
t.Fatalf("expected 1 merged box, got %d", len(result))
|
||
}
|
||
// The merged box's Bottom must be at least as large as any input Bottom.
|
||
// Known issue: see TODO in layout.go:236 and :284.
|
||
if result[0].Bottom < 300 {
|
||
t.Skipf("known issue: Bottom shrunk to %.1f (want >= 300) — deferred until pipeline alignment", result[0].Bottom)
|
||
}
|
||
}
|
||
|
||
func TestNaiveVerticalMerge(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, ColID: 0, X0: 50, X1: 550, Top: 100, Bottom: 112, Text: "第一段", LayoutNo: "1", LayoutType: "text"},
|
||
{PageNumber: 0, ColID: 0, X0: 50, X1: 550, Top: 114, Bottom: 126, Text: "续文", LayoutNo: "1", LayoutType: "text"},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 5}
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box, got %d: %v", len(result), result)
|
||
}
|
||
if len(result) > 0 && !strings.Contains(result[0].Text, "第一段") {
|
||
t.Errorf("merged text should contain '第一段': got %q", result[0].Text)
|
||
}
|
||
}
|
||
|
||
func TestNaiveVerticalMergeNonMerge(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, ColID: 0, X0: 50, X1: 550, Top: 100, Bottom: 112, Text: "第一段。", LayoutNo: "1", LayoutType: "text"},
|
||
{PageNumber: 0, ColID: 0, X0: 50, X1: 550, Top: 300, Bottom: 312, Text: "第二段。", LayoutNo: "1", LayoutType: "text"},
|
||
}
|
||
meanH := map[int]float64{0: 12}
|
||
meanW := map[int]float64{0: 5}
|
||
result := NaiveVerticalMerge(boxes, meanH, meanW, false)
|
||
if len(result) != 2 {
|
||
t.Errorf("expected 2 separate boxes (large gap), got %d", len(result))
|
||
}
|
||
}
|
||
|
||
// ---- 重构辅助函数的测试 ----
|
||
|
||
func TestGroupBoxesByPage(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 1, Text: "page1-box1"},
|
||
{PageNumber: 0, Text: "page0-box1"},
|
||
{PageNumber: 0, Text: "page0-box2"},
|
||
{PageNumber: 2, Text: "page2-box1"},
|
||
{PageNumber: 1, Text: "page1-box2"},
|
||
}
|
||
pageGroups, sortedPages := groupBoxesByPage(boxes)
|
||
|
||
if len(sortedPages) != 3 {
|
||
t.Errorf("expected 3 unique pages, got %d", len(sortedPages))
|
||
}
|
||
if sortedPages[0] != 0 || sortedPages[1] != 1 || sortedPages[2] != 2 {
|
||
t.Errorf("pages should be sorted [0,1,2], got %v", sortedPages)
|
||
}
|
||
if len(pageGroups[0]) != 2 {
|
||
t.Errorf("page 0 should have 2 boxes, got %d", len(pageGroups[0]))
|
||
}
|
||
if len(pageGroups[1]) != 2 {
|
||
t.Errorf("page 1 should have 2 boxes, got %d", len(pageGroups[1]))
|
||
}
|
||
if len(pageGroups[2]) != 1 {
|
||
t.Errorf("page 2 should have 1 box, got %d", len(pageGroups[2]))
|
||
}
|
||
if boxes[pageGroups[0][0]].Text != "page0-box1" {
|
||
t.Errorf("first page0 box index incorrect")
|
||
}
|
||
}
|
||
|
||
func TestGroupBoxesByPage_Empty(t *testing.T) {
|
||
pageGroups, sortedPages := groupBoxesByPage(nil)
|
||
if len(pageGroups) != 0 || len(sortedPages) != 0 {
|
||
t.Error("empty input should return empty result")
|
||
}
|
||
|
||
pageGroups, sortedPages = groupBoxesByPage([]pdf.TextBox{})
|
||
if len(pageGroups) != 0 || len(sortedPages) != 0 {
|
||
t.Error("empty input should return empty result")
|
||
}
|
||
}
|
||
|
||
func TestGroupBoxesByPage_SinglePage(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 5, Text: "box1"},
|
||
{PageNumber: 5, Text: "box2"},
|
||
}
|
||
pageGroups, sortedPages := groupBoxesByPage(boxes)
|
||
if len(sortedPages) != 1 || sortedPages[0] != 5 {
|
||
t.Errorf("expected single page 5, got %v", sortedPages)
|
||
}
|
||
if len(pageGroups[5]) != 2 {
|
||
t.Errorf("page 5 should have 2 boxes, got %d", len(pageGroups[5]))
|
||
}
|
||
}
|
||
|
||
func TestShouldMergeBoxes(t *testing.T) {
|
||
t.Run("should merge - basic case", func(t *testing.T) {
|
||
prev := &pdf.TextBox{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "前一句",
|
||
LayoutNo: "1",
|
||
}
|
||
curr := &pdf.TextBox{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: "后一句",
|
||
LayoutNo: "1",
|
||
}
|
||
if !shouldMergeBoxes(prev, curr, 12, 200, false) {
|
||
t.Error("should merge basic case")
|
||
}
|
||
})
|
||
|
||
t.Run("should NOT merge - different layoutNo", func(t *testing.T) {
|
||
prev := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250}
|
||
curr := &pdf.TextBox{PageNumber: 0, LayoutNo: "2", Top: 114, Bottom: 126, X0: 50, X1: 250}
|
||
if shouldMergeBoxes(prev, curr, 12, 200, false) {
|
||
t.Error("should not merge different layoutNo")
|
||
}
|
||
})
|
||
|
||
t.Run("should NOT merge - gap too large", func(t *testing.T) {
|
||
prev := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250}
|
||
curr := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 200, Bottom: 212, X0: 50, X1: 250}
|
||
if shouldMergeBoxes(prev, curr, 12, 200, false) {
|
||
t.Error("should not merge large gap")
|
||
}
|
||
})
|
||
|
||
t.Run("should NOT merge - overlap too small", func(t *testing.T) {
|
||
prev := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 100}
|
||
curr := &pdf.TextBox{PageNumber: 0, LayoutNo: "1", Top: 114, Bottom: 126, X0: 200, X1: 250}
|
||
if shouldMergeBoxes(prev, curr, 12, 200, false) {
|
||
t.Error("should not merge small overlap")
|
||
}
|
||
})
|
||
|
||
t.Run("should merge - comma override period anti", func(t *testing.T) {
|
||
prev := &pdf.TextBox{
|
||
PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250,
|
||
Text: "前一句。",
|
||
}
|
||
curr := &pdf.TextBox{
|
||
PageNumber: 0, LayoutNo: "1", Top: 114, Bottom: 126, X0: 50, X1: 250,
|
||
Text: ", 续句",
|
||
}
|
||
if !shouldMergeBoxes(prev, curr, 12, 200, false) {
|
||
t.Error("should merge when comma overrides period anti")
|
||
}
|
||
})
|
||
|
||
t.Run("should NOT merge - english period anti", func(t *testing.T) {
|
||
prev := &pdf.TextBox{
|
||
PageNumber: 0, LayoutNo: "1", Top: 100, Bottom: 112, X0: 50, X1: 250,
|
||
Text: "End of sentence.",
|
||
}
|
||
curr := &pdf.TextBox{
|
||
PageNumber: 0, LayoutNo: "1", Top: 114, Bottom: 126, X0: 50, X1: 250,
|
||
Text: "Next sentence",
|
||
}
|
||
if shouldMergeBoxes(prev, curr, 12, 200, true) {
|
||
t.Error("should not merge english period anti")
|
||
}
|
||
})
|
||
}
|
||
|
||
func TestMergeTwoBoxes(t *testing.T) {
|
||
prev := pdf.TextBox{
|
||
PageNumber: 0, X0: 50, X1: 200, Top: 100, Bottom: 112,
|
||
Text: "第一行",
|
||
LayoutNo: "1",
|
||
}
|
||
curr := pdf.TextBox{
|
||
PageNumber: 0, X0: 60, X1: 250, Top: 114, Bottom: 130,
|
||
Text: "第二行",
|
||
LayoutNo: "1",
|
||
}
|
||
|
||
result := mergeTwoBoxes(prev, curr)
|
||
|
||
expectedText := "第一行 第二行"
|
||
if result.Text != expectedText {
|
||
t.Errorf("expected text %q, got %q", expectedText, result.Text)
|
||
}
|
||
if result.X0 != 50 {
|
||
t.Errorf("expected X0 50, got %f", result.X0)
|
||
}
|
||
if result.X1 != 250 {
|
||
t.Errorf("expected X1 250, got %f", result.X1)
|
||
}
|
||
if result.Bottom != 130 {
|
||
t.Errorf("expected Bottom 130, got %f", result.Bottom)
|
||
}
|
||
if result.LayoutNo != "1" {
|
||
t.Errorf("expected LayoutNo preserved")
|
||
}
|
||
}
|
||
|
||
func TestMergeTwoBoxes_TrimWhitespace(t *testing.T) {
|
||
prev := pdf.TextBox{Text: " first line "}
|
||
curr := pdf.TextBox{Text: " second line "}
|
||
|
||
result := mergeTwoBoxes(prev, curr)
|
||
|
||
if result.Text != "first line second line" {
|
||
t.Errorf("text should be trimmed and joined, got %q", result.Text)
|
||
}
|
||
}
|
||
|
||
func TestProcessPageBoxes(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 114, Bottom: 126,
|
||
Text: "第二句",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "第一句",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
|
||
result := processPageBoxes(boxes, 12, 200, false)
|
||
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box, got %d", len(result))
|
||
}
|
||
if !strings.Contains(result[0].Text, "第一句") || !strings.Contains(result[0].Text, "第二句") {
|
||
t.Errorf("merged text should contain both parts, got %q", result[0].Text)
|
||
}
|
||
}
|
||
|
||
func TestProcessPageBoxes_WhitespaceBox(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "第一句",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 113, Bottom: 115,
|
||
Text: " ",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 116, Bottom: 128,
|
||
Text: "第二句",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
|
||
result := processPageBoxes(boxes, 12, 200, false)
|
||
|
||
if len(result) != 1 {
|
||
t.Errorf("expected 1 merged box, got %d", len(result))
|
||
}
|
||
}
|
||
|
||
func TestProcessPageBoxes_NoMerge(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 100, Bottom: 112,
|
||
Text: "第一句。",
|
||
LayoutNo: "1",
|
||
},
|
||
{
|
||
PageNumber: 0, X0: 50, X1: 250, Top: 200, Bottom: 212,
|
||
Text: "第二句",
|
||
LayoutNo: "1",
|
||
},
|
||
}
|
||
|
||
result := processPageBoxes(boxes, 12, 200, false)
|
||
|
||
if len(result) != 2 {
|
||
t.Errorf("expected 2 boxes, got %d", len(result))
|
||
}
|
||
}
|
||
|
||
// ── Column-assignment helper tests ──────────────────────────────────
|
||
|
||
func TestExtractX0Values(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 50, X1: 200},
|
||
{PageNumber: 0, X0: 30, X1: 100},
|
||
{PageNumber: 0, X0: 80, X1: 300},
|
||
}
|
||
x0s, minX0, maxX1 := extractX0Values(boxes, []int{0, 1, 2})
|
||
if len(x0s) != 3 {
|
||
t.Fatalf("expected 3 x0s, got %d", len(x0s))
|
||
}
|
||
if x0s[0] != 50 || x0s[1] != 30 || x0s[2] != 80 {
|
||
t.Errorf("x0s mismatch: %v", x0s)
|
||
}
|
||
if minX0 != 30 {
|
||
t.Errorf("minX0 = %v, want 30", minX0)
|
||
}
|
||
if maxX1 != 300 {
|
||
t.Errorf("maxX1 = %v, want 300", maxX1)
|
||
}
|
||
}
|
||
|
||
func TestApplyIndentTolerance(t *testing.T) {
|
||
values := []float64{100, 105, 200, 210}
|
||
applyIndentTolerance(values, 100, 10)
|
||
if values[0] != 100 || values[1] != 100 {
|
||
t.Errorf("close x0s should be adjusted to minX0: %v", values)
|
||
}
|
||
if values[2] != 200 || values[3] != 210 {
|
||
t.Errorf("distant x0s should remain unchanged: %v", values)
|
||
}
|
||
}
|
||
|
||
func TestApplyIndentTolerance_Zero(t *testing.T) {
|
||
values := []float64{100, 101, 200}
|
||
applyIndentTolerance(values, 100, 0)
|
||
if values[1] != 101 {
|
||
t.Errorf("zero tolerance: x0s should be unchanged, got %v", values)
|
||
}
|
||
}
|
||
|
||
func TestApplyIndentTolerance_Negative(t *testing.T) {
|
||
values := []float64{-100, -95, 0, 50}
|
||
applyIndentTolerance(values, -100, 10)
|
||
if values[0] != -100 || values[1] != -100 {
|
||
t.Errorf("negative x0s close to minX0 should be adjusted: %v", values)
|
||
}
|
||
}
|
||
|
||
func TestFindBestK_SingleCluster(t *testing.T) {
|
||
// Note: KMeans1D uses random initialization, so non-identical values
|
||
// may occasionally produce k>1. This test verifies the function runs
|
||
// without error and returns k>=1 (not a correctness check).
|
||
x0s := []float64{100, 99, 101}
|
||
bestK, _ := findBestK(x0s, len(x0s))
|
||
if bestK < 1 {
|
||
t.Errorf("expected bestK>=1, got %d", bestK)
|
||
}
|
||
}
|
||
|
||
func TestFindBestK_TwoColumns(t *testing.T) {
|
||
x0s := []float64{50, 55, 60, 200, 210, 220}
|
||
bestK, _ := findBestK(x0s, len(x0s))
|
||
if bestK != 2 {
|
||
t.Errorf("two columns: expected bestK=2, got %d", bestK)
|
||
}
|
||
}
|
||
|
||
func TestFindBestK_OneValue(t *testing.T) {
|
||
x0s := []float64{100}
|
||
bestK, _ := findBestK(x0s, len(x0s))
|
||
if bestK != 1 {
|
||
t.Errorf("single value: expected bestK=1, got %d", bestK)
|
||
}
|
||
}
|
||
|
||
func TestFindBestK_Identical(t *testing.T) {
|
||
x0s := []float64{100, 100, 100, 100, 100}
|
||
bestK, _ := findBestK(x0s, len(x0s))
|
||
if bestK != 1 {
|
||
t.Errorf("identical values: expected bestK=1, got %d", bestK)
|
||
}
|
||
}
|
||
|
||
func TestRemapLabelsByCentroidOrder_Ordered(t *testing.T) {
|
||
centroids := []float64{50, 200, 400}
|
||
remap := remapLabelsByCentroidOrder(centroids)
|
||
if remap[0] != 0 || remap[1] != 1 || remap[2] != 2 {
|
||
t.Errorf("ordered centroids: expected 0->0,1->1,2->2, got %v", remap)
|
||
}
|
||
}
|
||
|
||
func TestRemapLabelsByCentroidOrder_Unordered(t *testing.T) {
|
||
centroids := []float64{200, 50, 400}
|
||
remap := remapLabelsByCentroidOrder(centroids)
|
||
if remap[0] != 1 || remap[1] != 0 || remap[2] != 2 {
|
||
t.Errorf("unordered centroids: expected {0:1,1:0,2:2}, got %v", remap)
|
||
}
|
||
}
|
||
|
||
func TestRemapLabelsByCentroidOrder_Nil(t *testing.T) {
|
||
remap := remapLabelsByCentroidOrder(nil)
|
||
if len(remap) != 0 {
|
||
t.Errorf("nil centroids: expected empty map, got %v", remap)
|
||
}
|
||
}
|
||
|
||
func TestDetermineBestKForPage_SingleBox(t *testing.T) {
|
||
boxes := []pdf.TextBox{{PageNumber: 0, X0: 100, X1: 200}}
|
||
result := make([]pdf.TextBox, len(boxes))
|
||
copy(result, boxes)
|
||
pageCols := make(map[int]int)
|
||
determineBestKForPage(boxes, result, []int{0}, 0, pageCols)
|
||
if pageCols[0] != 1 {
|
||
t.Errorf("single box: expected pageCols[0]=1, got %d", pageCols[0])
|
||
}
|
||
if result[0].ColID != 0 {
|
||
t.Errorf("single box: expected ColID=0, got %d", result[0].ColID)
|
||
}
|
||
}
|
||
|
||
func TestDetermineBestKForPage_TwoColumns(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 50, X1: 100},
|
||
{PageNumber: 0, X0: 55, X1: 100},
|
||
{PageNumber: 0, X0: 300, X1: 400},
|
||
{PageNumber: 0, X0: 310, X1: 400},
|
||
}
|
||
result := make([]pdf.TextBox, len(boxes))
|
||
copy(result, boxes)
|
||
pageCols := make(map[int]int)
|
||
determineBestKForPage(boxes, result, []int{0, 1, 2, 3}, 0, pageCols)
|
||
if pageCols[0] != 2 {
|
||
t.Errorf("two distinct columns: expected pageCols[0]=2, got %d", pageCols[0])
|
||
}
|
||
}
|
||
|
||
func TestAssignmentHelpers_IndentTolerance(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 50, X1: 150, Top: 10, Bottom: 30},
|
||
{PageNumber: 0, X0: 205, X1: 350, Top: 10, Bottom: 30},
|
||
{PageNumber: 0, X0: 58, X1: 150, Top: 40, Bottom: 60},
|
||
}
|
||
result := make([]pdf.TextBox, len(boxes))
|
||
copy(result, boxes)
|
||
pageCols := make(map[int]int)
|
||
determineBestKForPage(boxes, result, []int{0, 1, 2}, 0, pageCols)
|
||
if pageCols[0] != 2 {
|
||
t.Errorf("expected 2 columns after indent tolerance, got %d", pageCols[0])
|
||
}
|
||
}
|
||
|
||
func TestAssignColIDsForPage_Normal(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 50, X1: 100},
|
||
{PageNumber: 0, X0: 200, X1: 300},
|
||
}
|
||
result := make([]pdf.TextBox, len(boxes))
|
||
copy(result, boxes)
|
||
pageCols := map[int]int{0: 2}
|
||
assignColIDsForPage(boxes, result, []int{0, 1}, 0, pageCols)
|
||
if result[0].ColID != 0 || result[1].ColID != 1 {
|
||
t.Errorf("expected ColIDs 0,1 but got %d,%d", result[0].ColID, result[1].ColID)
|
||
}
|
||
}
|
||
|
||
func TestAssignColIDsForPage_KTooLarge(t *testing.T) {
|
||
boxes := []pdf.TextBox{
|
||
{PageNumber: 0, X0: 100, X1: 200},
|
||
}
|
||
result := make([]pdf.TextBox, len(boxes))
|
||
copy(result, boxes)
|
||
pageCols := map[int]int{0: 3}
|
||
assignColIDsForPage(boxes, result, []int{0}, 0, pageCols)
|
||
if result[0].ColID != 0 {
|
||
t.Errorf("expected ColID=0 (k clamped to 1), got %d", result[0].ColID)
|
||
}
|
||
}
|