Agent-System/internal/mineru/client_test.go

package mineru_test

import (
	"archive/zip"
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"chat-agent-system/internal/mineru"
)

func TestMarkdownFromContentList_PageAnchors(t *testing.T) {
	raw, _ := json.Marshal([]map[string]any{
		{"type": "text", "text": "第一页内容", "page_idx": 0, "text_level": 0},
		{"type": "title", "text": "第二章", "page_idx": 1, "text_level": 2},
		{"type": "text", "text": "第二页正文", "page_idx": 1, "text_level": 0},
	})
	md, pages, err := mineru.MarkdownFromContentList(raw)
	if err != nil {
		t.Fatal(err)
	}
	if pages != 2 {
		t.Fatalf("pages=%d", pages)
	}
	if !strings.Contains(md, "<!-- page 1 -->") || !strings.Contains(md, "<!-- page 2 -->") {
		t.Fatalf("missing page anchors:\n%s", md)
	}
	if !strings.Contains(md, "## 第二章") {
		t.Fatalf("missing heading:\n%s", md)
	}
}

func TestExtractFromZip_PrefersContentList(t *testing.T) {
	var buf bytes.Buffer
	zw := zip.NewWriter(&buf)
	w, _ := zw.Create("full.md")
	_, _ = w.Write([]byte("# stale full\n"))
	list, _ := json.Marshal([]map[string]any{
		{"type": "text", "text": "zip内容", "page_idx": 0},
	})
	w2, _ := zw.Create("demo_content_list.json")
	_, _ = w2.Write(list)
	_ = zw.Close()

	md, _, pages, err := mineru.ExtractFromZip(buf.Bytes())
	if err != nil {
		t.Fatal(err)
	}
	if pages != 1 || !strings.Contains(md, "zip内容") {
		t.Fatalf("md=%q pages=%d", md, pages)
	}
	if strings.Contains(md, "stale full") {
		t.Fatal("should prefer content_list over full.md")
	}
}

func TestClient_ParseFile_MockServer(t *testing.T) {
	var pollCount atomic.Int32
	var zipBytes bytes.Buffer
	zw := zip.NewWriter(&zipBytes)
	list, _ := json.Marshal([]map[string]any{
		{"type": "text", "text": "hello mineru", "page_idx": 0},
	})
	w, _ := zw.Create("full.md")
	_, _ = w.Write([]byte("# fallback\n"))
	w2, _ := zw.Create("x_content_list.json")
	_, _ = w2.Write(list)
	_ = zw.Close()

	var baseURL string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.URL.Path == "/upload" && r.Method == http.MethodPut:
			w.WriteHeader(200)
		case r.URL.Path == "/result.zip":
			_, _ = w.Write(zipBytes.Bytes())
		case r.URL.Path == "/api/v4/file-urls/batch":
			_ = json.NewEncoder(w).Encode(map[string]any{
				"code": 0,
				"data": map[string]any{
					"batch_id":  "batch-1",
					"file_urls": []string{baseURL + "/upload"},
				},
			})
		case strings.HasPrefix(r.URL.Path, "/api/v4/extract-results/batch/"):
			n := pollCount.Add(1)
			state := "running"
			zipURL := ""
			if n >= 2 {
				state = "done"
				zipURL = baseURL + "/result.zip"
			}
			_ = json.NewEncoder(w).Encode(map[string]any{
				"code": 0,
				"data": map[string]any{
					"extract_result": []map[string]any{{
						"file_name":    "doc.pdf",
						"state":        state,
						"full_zip_url": zipURL,
					}},
				},
			})
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()
	baseURL = srv.URL

	client := mineru.NewClient(mineru.Config{
		Endpoint:    srv.URL,
		APIKey:      "test-token",
		HTTPTimeout: 30 * time.Second,
	})

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	res, err := client.ParseFile(ctx, "doc.pdf", []byte("%PDF-1.4 mock"))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(res.Markdown, "hello mineru") {
		t.Fatalf("markdown=%q", res.Markdown)
	}
}

func TestSupportsExt(t *testing.T) {
	if !mineru.SupportsExt("a.PDF") || mineru.SupportsExt("a.txt") {
		t.Fatal("ext filter")
	}
}