SnapGo/internal/infrastructure/ocr/client_test.go

package ocr

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/mmmy/snapgo/internal/domain"
)

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req)
}

func TestAliyunRecognizeTextSignsRawPNGRequest(t *testing.T) {
	var requestBody []byte
	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		if r.URL.Host != "ocr-api.cn-hangzhou.aliyuncs.com" {
			t.Fatalf("unexpected host %s", r.URL.Host)
		}
		if got := r.Header.Get("x-acs-action"); got != "RecognizeGeneral" {
			t.Fatalf("unexpected action %q", got)
		}
		if got := r.Header.Get("x-acs-version"); got != "2021-07-07" {
			t.Fatalf("unexpected version %q", got)
		}
		if got := r.Header.Get("x-acs-date"); got != "2026-07-08T01:02:03Z" {
			t.Fatalf("unexpected date %q", got)
		}
		if got := r.Header.Get("Authorization"); !strings.Contains(got, "ACS3-HMAC-SHA256 Credential=ak") {
			t.Fatalf("unexpected auth header %q", got)
		}
		var err error
		requestBody, err = io.ReadAll(r.Body)
		if err != nil {
			t.Fatalf("read body: %v", err)
		}
		return &http.Response{
			StatusCode: http.StatusOK,
			Header:     http.Header{"Content-Type": []string{"application/json"}},
			Body: io.NopCloser(bytes.NewBufferString(
				`{"Data":"{\"content\":\"第一行\\n第二行\"}"}`,
			)),
		}, nil
	})

	client, err := NewClient(domain.OCRProviderAliyun, domain.OCRProviderConfig{
		Endpoint:        "https://ocr-api.cn-hangzhou.aliyuncs.com",
		AccessKeyID:     "ak",
		AccessKeySecret: "sk",
		Action:          "RecognizeGeneral",
		Version:         "2021-07-07",
	})
	if err != nil {
		t.Fatalf("new client: %v", err)
	}
	client.httpClient.Transport = transport
	client.now = func() time.Time { return time.Date(2026, 7, 8, 1, 2, 3, 0, time.UTC) }
	client.nonce = func() string { return "nonce" }

	text, err := client.RecognizeText(context.Background(), []byte("png"))
	if err != nil {
		t.Fatalf("recognize text: %v", err)
	}
	if string(requestBody) != "png" {
		t.Fatalf("expected raw png body, got %q", string(requestBody))
	}
	if text != "第一行\n第二行" {
		t.Fatalf("unexpected OCR text %q", text)
	}
}

func TestVolcengineRecognizeTextSignsBase64JSONRequest(t *testing.T) {
	var requestBody map[string]string
	transport := roundTripFunc(func(r *http.Request) (*http.Response, error) {
		if got := r.URL.Query().Get("Action"); got != "OCRNormal" {
			t.Fatalf("unexpected action %q", got)
		}
		if got := r.URL.Query().Get("Version"); got != "2020-08-26" {
			t.Fatalf("unexpected version %q", got)
		}
		if got := r.Header.Get("X-Date"); got != "20260708T010203Z" {
			t.Fatalf("unexpected x-date %q", got)
		}
		if got := r.Header.Get("Authorization"); !strings.Contains(got, "HMAC-SHA256 Credential=ak/20260708/cn-north-1/cv/request") {
			t.Fatalf("unexpected auth header %q", got)
		}
		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		return &http.Response{
			StatusCode: http.StatusOK,
			Header:     http.Header{"Content-Type": []string{"application/json"}},
			Body: io.NopCloser(bytes.NewBufferString(
				`{"ResponseMetadata":{"RequestId":"r"},"Result":{"LineTexts":["你好","世界"]}}`,
			)),
		}, nil
	})

	client, err := NewClient(domain.OCRProviderVolcengine, domain.OCRProviderConfig{
		Endpoint:        "https://visual.volcengineapi.com",
		AccessKeyID:     "ak",
		AccessKeySecret: "sk",
		Region:          "cn-north-1",
		Service:         "cv",
		Action:          "OCRNormal",
		Version:         "2020-08-26",
	})
	if err != nil {
		t.Fatalf("new client: %v", err)
	}
	client.httpClient.Transport = transport
	client.now = func() time.Time { return time.Date(2026, 7, 8, 1, 2, 3, 0, time.UTC) }

	text, err := client.RecognizeText(context.Background(), []byte("png"))
	if err != nil {
		t.Fatalf("recognize text: %v", err)
	}
	if requestBody["image_base64"] != base64.StdEncoding.EncodeToString([]byte("png")) {
		t.Fatalf("expected base64 image body, got %#v", requestBody)
	}
	if text != "你好\n世界" {
		t.Fatalf("unexpected OCR text %q", text)
	}
}

func TestParseOCRTextSupportsWordsResult(t *testing.T) {
	text, err := parseOCRText([]byte(`{"data":{"words_result":[{"words":"foo"},{"words":"bar"}]}}`))
	if err != nil {
		t.Fatalf("parse ocr text: %v", err)
	}
	if text != "foo\nbar" {
		t.Fatalf("unexpected OCR text %q", text)
	}
}

func TestParseOCRTextPrefersOverallContentOverBlocks(t *testing.T) {
	text, err := parseOCRText([]byte(`{
		"Data": "{\"content\":\"整体识别文本\",\"prism_wordsInfo\":[{\"word\":\"整体\"},{\"word\":\"识别\"},{\"word\":\"文本\"}]}"
	}`))
	if err != nil {
		t.Fatalf("parse ocr text: %v", err)
	}
	if text != "整体识别文本" {
		t.Fatalf("expected overall content only, got %q", text)
	}
}

func TestParseOCRTextPrefersResultTextOverWordsResult(t *testing.T) {
	text, err := parseOCRText([]byte(`{
		"Result": {
			"text": "hello world",
			"words_result": [{"words":"hello"},{"words":"world"}]
		}
	}`))
	if err != nil {
		t.Fatalf("parse ocr text: %v", err)
	}
	if text != "hello world" {
		t.Fatalf("expected result text only, got %q", text)
	}
}

func TestParseOCRTextReturnsProviderError(t *testing.T) {
	_, err := parseOCRText([]byte(`{"ResponseMetadata":{"Error":{"Code":"BadRequest","Message":"bad image"}}}`))
	if err == nil || !strings.Contains(err.Error(), "bad image") {
		t.Fatalf("expected provider error, got %v", err)
	}
}