feat: add screenshot OCR extraction

2026-07-08 00:42:49 +08:00
parent f90612976f
commit baec01610c
14 changed files with 1420 additions and 46 deletions
@@ -0,0 +1,150 @@
+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 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)
+	}
+}