ragflow/internal/deepdoc/parser/pdf/table_rotate_test.go

package parser

import (
	"context"
	"image"
	"testing"
)

// mockRotationDoc implements DocAnalyzer with deterministic OCR results per angle.
// The mock tracks the call sequence: evaluateTableOrientation tests angles in
// order 0°, 90°, 180°, 270°. Each call to OCRDetect increments an internal
// counter and returns data for the corresponding angle.
type mockRotationDoc struct {
	// angle → {regions count, average confidence, error}
	angles map[int]struct {
		regions int
		avgConf float64
		err     error
	}
	callSeq int // incremented per OCRDetect call, selects the angle's data
}

var rotationOrder = []int{0, 90, 180, 270}

func (m *mockRotationDoc) DLA(_ context.Context, _ image.Image) ([]DLARegion, error) { return nil, nil }
func (m *mockRotationDoc) TSR(_ context.Context, _ image.Image) ([]TSRCell, error)   { return nil, nil }
func (m *mockRotationDoc) OCR(_ image.Image) (string, error)                         { return "", nil }
func (m *mockRotationDoc) Health() bool                                              { return true }
func (m *mockRotationDoc) ModelType() ModelType                                      { return ModelSaas }

func (m *mockRotationDoc) currentAngle() int {
	idx := m.callSeq % len(rotationOrder)
	return rotationOrder[idx]
}

func (m *mockRotationDoc) OCRDetect(_ context.Context, img image.Image) ([]OCRBox, error) {
	defer func() { m.callSeq++ }()
	angle := m.currentAngle()
	cfg, ok := m.angles[angle]
	if !ok {
		cfg = m.angles[0] // fallback to 0° config
	}
	if cfg.err != nil {
		return nil, cfg.err
	}
	if cfg.regions == 0 {
		return nil, nil
	}
	w, h := img.Bounds().Dx(), img.Bounds().Dy()
	boxes := make([]OCRBox, cfg.regions)
	step := w / (cfg.regions + 1)
	for i := 0; i < cfg.regions; i++ {
		x := step * (i + 1)
		boxes[i] = OCRBox{
			X0: float64(x), Y0: float64(h / 4),
			X1: float64(x + 20), Y1: float64(h / 4),
			X2: float64(x + 20), Y2: float64(h * 3 / 4),
			X3: float64(x), Y3: float64(h * 3 / 4),
		}
	}
	return boxes, nil
}

func (m *mockRotationDoc) OCRRecognizeBatch(_ context.Context, cropped []image.Image) ([][]OCRText, []error) {
	results := make([][]OCRText, len(cropped))
	errs := make([]error, len(cropped))
	for i, img := range cropped {
		results[i], errs[i] = m.OCRRecognize(context.Background(), img)
	}
	return results, errs
}

func (m *mockRotationDoc) OCRRecognize(_ context.Context, _ image.Image) ([]OCRText, error) {
	angle := rotationOrder[(m.callSeq-1)%len(rotationOrder)] // use angle from last Detect call
	cfg, ok := m.angles[angle]
	if !ok {
		cfg = m.angles[0]
	}
	if cfg.err != nil {
		return nil, cfg.err
	}
	if cfg.regions == 0 {
		return nil, nil
	}
	texts := make([]OCRText, cfg.regions)
	for i := 0; i < cfg.regions; i++ {
		texts[i] = OCRText{Text: "X", Confidence: cfg.avgConf}
	}
	return texts, nil
}

func makeTestTableImage() image.Image {
	return image.NewRGBA(image.Rect(0, 0, 200, 100))
}

func TestEvaluateTableOrientation(t *testing.T) {
	t.Run("normal table 0° wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0: {regions: 10, avgConf: 0.9},
			},
		}
		angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("90° rotated table wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {regions: 2, avgConf: 0.2},
				90:  {regions: 10, avgConf: 0.9},
				180: {regions: 2, avgConf: 0.2},
				270: {regions: 2, avgConf: 0.2},
			},
		}
		angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 90 {
			t.Errorf("expected 90°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("180° rotated table wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {regions: 1, avgConf: 0.1},
				90:  {regions: 1, avgConf: 0.1},
				180: {regions: 8, avgConf: 0.85},
				270: {regions: 1, avgConf: 0.1},
			},
		}
		angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 180 {
			t.Errorf("expected 180°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("270° rotated table wins", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {regions: 1, avgConf: 0.1},
				90:  {regions: 1, avgConf: 0.1},
				180: {regions: 1, avgConf: 0.1},
				270: {regions: 9, avgConf: 0.88},
			},
		}
		angle, _, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 270 {
			t.Errorf("expected 270°, got %d° (scores: %v)", angle, scores)
		}
	})

	t.Run("threshold protection — 0° keeps when diff too small", func(t *testing.T) {
		// Region-count scoring: 8 vs 9 is too close (< 1.4×) → 0° wins.
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 8},
				90: {regions: 9},
			},
		}
		angle, _, _ := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0° (threshold protection), got %d°", angle)
		}
	})

	t.Run("threshold pass — 90° wins when region count is clearly higher", func(t *testing.T) {
		// 0° has few regions AND 90° has ≥1.4× more → 90° wins.
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:  {regions: 4},
				90: {regions: 10},
			},
		}
		angle, _, _ := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 90 {
			t.Errorf("expected 90° (threshold passed), got %d°", angle)
		}
	})

	t.Run("all angles fail OCR → fallback 0°", func(t *testing.T) {
		doc := &mockRotationDoc{
			angles: map[int]struct {
				regions int
				avgConf float64
				err     error
			}{
				0:   {err: errMockOCR},
				90:  {err: errMockOCR},
				180: {err: errMockOCR},
				270: {err: errMockOCR},
			},
		}
		angle, img, scores := evaluateTableOrientation(context.Background(), makeTestTableImage(), doc)
		if angle != 0 {
			t.Errorf("expected 0° fallback, got %d°", angle)
		}
		if img == nil {
			t.Error("expected non-nil fallback image")
		}
		for _, s := range scores {
			if s != 0 {
				t.Error("all scores should be 0 on OCR failure")
			}
		}
	})
}

var errMockOCR = &mockError{"mock OCR failure"}

type mockError struct{ msg string }

func (e *mockError) Error() string { return e.msg }