08c83e09e6
- Modified Schema methods across various providers (chunking, embedding, extractor, fetcher, converters) to accept an additional locale parameter, enhancing localization support. - Updated corresponding test cases to reflect the changes in method signatures and ensure proper functionality with the new locale argument.
407 lines
12 KiB
Go
407 lines
12 KiB
Go
package converters
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/yaoapp/yao/config"
|
|
kbtypes "github.com/yaoapp/yao/kb/types"
|
|
"github.com/yaoapp/yao/test"
|
|
)
|
|
|
|
func TestOCR_Make(t *testing.T) {
|
|
// Setup
|
|
test.Prepare(&testing.T{}, config.Conf)
|
|
defer test.Clean()
|
|
|
|
ocr := &OCR{}
|
|
|
|
t.Run("nil option should return error for missing vision converter", func(t *testing.T) {
|
|
_, err := ocr.Make(nil)
|
|
if err == nil {
|
|
t.Error("Expected error for missing vision converter")
|
|
}
|
|
if err.Error() != "vision converter is required for OCR processing" {
|
|
t.Errorf("Expected specific error message, got: %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("empty option should return error for missing vision converter", func(t *testing.T) {
|
|
option := &kbtypes.ProviderOption{}
|
|
_, err := ocr.Make(option)
|
|
if err == nil {
|
|
t.Error("Expected error for missing vision converter")
|
|
}
|
|
if err.Error() != "vision converter is required for OCR processing" {
|
|
t.Errorf("Expected specific error message, got: %v", err)
|
|
}
|
|
})
|
|
|
|
t.Run("should use global PDF configuration as defaults", func(t *testing.T) {
|
|
// Set up global PDF configuration
|
|
globalPDFConfig := &kbtypes.PDFConfig{
|
|
ConvertTool: "mutool",
|
|
ToolPath: "/usr/local/bin/mutool",
|
|
}
|
|
kbtypes.SetGlobalPDF(globalPDFConfig)
|
|
|
|
// Clean up after test
|
|
defer kbtypes.SetGlobalPDF(nil)
|
|
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
// This will fail because vision converter factory isn't set up in tests
|
|
// but we can verify the error shows the global config was used
|
|
_, err := ocr.Make(option)
|
|
if err == nil {
|
|
t.Error("Expected error due to mock factory limitation")
|
|
}
|
|
// In real usage with proper factory setup, this would work
|
|
// and would use mutool as the PDF tool and /usr/local/bin/mutool as the path
|
|
})
|
|
|
|
t.Run("properties should override global PDF configuration", func(t *testing.T) {
|
|
// Set up global PDF configuration
|
|
globalPDFConfig := &kbtypes.PDFConfig{
|
|
ConvertTool: "mutool",
|
|
ToolPath: "/usr/local/bin/mutool",
|
|
}
|
|
kbtypes.SetGlobalPDF(globalPDFConfig)
|
|
|
|
// Clean up after test
|
|
defer kbtypes.SetGlobalPDF(nil)
|
|
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"pdf_tool": "pdftoppm", // Override global mutool with pdftoppm
|
|
"pdf_tool_path": "/usr/bin/pdftoppm", // Override global path
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
// This will fail because vision converter factory isn't set up in tests
|
|
// but the properties would override the global configuration
|
|
_, err := ocr.Make(option)
|
|
if err == nil {
|
|
t.Error("Expected error due to mock factory limitation")
|
|
}
|
|
// In real usage, this would use pdftoppm instead of the global mutool setting
|
|
})
|
|
|
|
t.Run("option with OCR properties should set all values", func(t *testing.T) {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"mode": "concurrent",
|
|
"max_concurrency": 8,
|
|
"compress_size": 1024,
|
|
"force_image_mode": true,
|
|
"pdf_tool": "mutool",
|
|
"pdf_tool_path": "/usr/bin/mutool",
|
|
"pdf_dpi": 200,
|
|
"pdf_format": "jpg",
|
|
"pdf_quality": 85,
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// This will fail because vision converter factory isn't set up in tests
|
|
_, err := ocr.Make(option)
|
|
if err == nil {
|
|
t.Error("Expected error due to mock factory limitation")
|
|
}
|
|
// In real usage, this would work with proper factory setup
|
|
})
|
|
|
|
t.Run("should work without global PDF configuration", func(t *testing.T) {
|
|
// Ensure no global PDF configuration is set
|
|
kbtypes.SetGlobalPDF(nil)
|
|
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"pdf_tool": "pdftoppm",
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
// This will fail because vision converter factory isn't set up in tests
|
|
_, err := ocr.Make(option)
|
|
if err == nil {
|
|
t.Error("Expected error due to mock factory limitation")
|
|
}
|
|
// In real usage, this would work and use hardcoded defaults for unspecified PDF settings
|
|
})
|
|
|
|
t.Run("mode selection should work correctly", func(t *testing.T) {
|
|
testCases := []struct {
|
|
mode string
|
|
shouldWork bool
|
|
}{
|
|
{"queue", true},
|
|
{"concurrent", true},
|
|
{"invalid", true}, // Should default to queue mode
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"mode": tc.mode,
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// This will fail due to factory setup, but we're testing the parsing logic
|
|
_, err := ocr.Make(option)
|
|
// We expect error due to vision converter factory not being set up
|
|
if err == nil {
|
|
t.Errorf("Expected error for mode %s due to test limitations", tc.mode)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("PDF tool selection should work correctly", func(t *testing.T) {
|
|
testCases := []struct {
|
|
tool string
|
|
}{
|
|
{"pdftoppm"},
|
|
{"mutool"},
|
|
{"imagemagick"},
|
|
{"convert"},
|
|
{"invalid"}, // Should default
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"pdf_tool": tc.tool,
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// This will fail due to factory setup, but we're testing the parsing logic
|
|
_, err := ocr.Make(option)
|
|
// We expect error due to vision converter factory not being set up
|
|
if err == nil {
|
|
t.Errorf("Expected error for PDF tool %s due to test limitations", tc.tool)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("numeric values should handle both int and float64", func(t *testing.T) {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"max_concurrency": 8, // int
|
|
"compress_size": 512.0, // float64
|
|
"pdf_dpi": 150.0, // float64 -> int
|
|
"pdf_quality": 90, // int
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// This will fail due to factory setup, but we're testing the parsing logic
|
|
_, err := ocr.Make(option)
|
|
// We expect error due to vision converter factory not being set up
|
|
if err == nil {
|
|
t.Error("Expected error due to test limitations")
|
|
}
|
|
})
|
|
|
|
t.Run("boolean values should be handled correctly", func(t *testing.T) {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"force_image_mode": true,
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// This will fail due to factory setup, but we're testing the parsing logic
|
|
_, err := ocr.Make(option)
|
|
// We expect error due to vision converter factory not being set up
|
|
if err == nil {
|
|
t.Error("Expected error due to test limitations")
|
|
}
|
|
})
|
|
|
|
t.Run("invalid property types should be ignored", func(t *testing.T) {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"mode": 123, // invalid type
|
|
"max_concurrency": "invalid", // invalid type
|
|
"compress_size": "invalid", // invalid type
|
|
"force_image_mode": "invalid", // invalid type
|
|
"pdf_dpi": "invalid", // invalid type
|
|
"vision": map[string]interface{}{
|
|
"converter": "__yao.vision",
|
|
"properties": map[string]interface{}{
|
|
"connector": "openai.gpt-4o-mini",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
// This will fail due to factory setup, but we're testing the parsing logic
|
|
_, err := ocr.Make(option)
|
|
// We expect error due to vision converter factory not being set up
|
|
if err == nil {
|
|
t.Error("Expected error due to test limitations")
|
|
}
|
|
})
|
|
|
|
t.Run("invalid vision converter should return error", func(t *testing.T) {
|
|
option := &kbtypes.ProviderOption{
|
|
Properties: map[string]interface{}{
|
|
"vision": "invalid_format", // should be a map
|
|
},
|
|
}
|
|
_, err := ocr.Make(option)
|
|
if err == nil {
|
|
t.Error("Expected error for invalid vision converter format")
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestOCR_AutoDetect(t *testing.T) {
|
|
ocr := &OCR{
|
|
Autodetect: []string{".pdf", ".jpg", ".png", ".gif", "application/pdf", "image/jpeg"},
|
|
MatchPriority: 10,
|
|
}
|
|
|
|
t.Run("should detect .pdf files", func(t *testing.T) {
|
|
match, priority, err := ocr.AutoDetect("document.pdf", "")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if !match {
|
|
t.Error("Expected match for .pdf file")
|
|
}
|
|
if priority != 10 {
|
|
t.Errorf("Expected priority 10, got %d", priority)
|
|
}
|
|
})
|
|
|
|
t.Run("should detect .jpg files", func(t *testing.T) {
|
|
match, priority, err := ocr.AutoDetect("scan.jpg", "")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if !match {
|
|
t.Error("Expected match for .jpg file")
|
|
}
|
|
if priority != 10 {
|
|
t.Errorf("Expected priority 10, got %d", priority)
|
|
}
|
|
})
|
|
|
|
t.Run("should detect .png files", func(t *testing.T) {
|
|
match, priority, err := ocr.AutoDetect("screenshot.png", "")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if !match {
|
|
t.Error("Expected match for .png file")
|
|
}
|
|
if priority != 10 {
|
|
t.Errorf("Expected priority 10, got %d", priority)
|
|
}
|
|
})
|
|
|
|
t.Run("should detect by content type", func(t *testing.T) {
|
|
match, priority, err := ocr.AutoDetect("unknown", "application/pdf")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if !match {
|
|
t.Error("Expected match for application/pdf content type")
|
|
}
|
|
if priority != 10 {
|
|
t.Errorf("Expected priority 10, got %d", priority)
|
|
}
|
|
})
|
|
|
|
t.Run("should detect image content types", func(t *testing.T) {
|
|
match, priority, err := ocr.AutoDetect("unknown", "image/jpeg")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if !match {
|
|
t.Error("Expected match for image/jpeg content type")
|
|
}
|
|
if priority != 10 {
|
|
t.Errorf("Expected priority 10, got %d", priority)
|
|
}
|
|
})
|
|
|
|
t.Run("should not detect unsupported files", func(t *testing.T) {
|
|
match, priority, err := ocr.AutoDetect("video.mp4", "video/mp4")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if match {
|
|
t.Error("Expected no match for .mp4 file")
|
|
}
|
|
if priority != 0 {
|
|
t.Errorf("Expected priority 0, got %d", priority)
|
|
}
|
|
})
|
|
|
|
t.Run("empty autodetect should not match", func(t *testing.T) {
|
|
emptyOCR := &OCR{}
|
|
match, priority, err := emptyOCR.AutoDetect("document.pdf", "application/pdf")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if match {
|
|
t.Error("Expected no match when autodetect is empty")
|
|
}
|
|
if priority != 0 {
|
|
t.Errorf("Expected priority 0, got %d", priority)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestOCR_Schema(t *testing.T) {
|
|
ocr := &OCR{}
|
|
schema, err := ocr.Schema(nil, "en")
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got %v", err)
|
|
}
|
|
if schema != nil {
|
|
t.Error("Expected nil schema")
|
|
}
|
|
}
|