ragflow/internal/parser/parser/markdown_parser_test.go

package parser

import (
	"encoding/base64"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestMarkdownParser_ParseWithResult_Basic(t *testing.T) {
	ctx := t.Context()
	p, err := NewMarkdownParser(GoMarkdown)
	if err != nil {
		t.Fatalf("NewMarkdownParser: %v", err)
	}
	md := "# Hello\n\nThis is a paragraph.\n\n* List item 1\n* List item 2\n\n```go\nfunc main() {}\n```\n"
	res := p.ParseWithResult(ctx, "test.md", []byte(md))
	if res.Err != nil {
		t.Fatalf("ParseWithResult: %v", res.Err)
	}
	if got, want := res.OutputFormat, "json"; got != want {
		t.Fatalf("OutputFormat = %q, want %q", got, want)
	}
	if len(res.JSON) == 0 {
		t.Fatal("JSON is empty; want at least one item")
	}
	// Verify heading
	if got, _ := res.JSON[0]["text"].(string); got != "Hello" {
		t.Fatalf("first item text = %q, want %q", got, "Hello")
	}
	if got, _ := res.JSON[0]["ck_type"].(string); got != "heading" {
		t.Fatalf("first item ck_type = %q, want %q", got, "heading")
	}
}

func TestMarkdownParser_ParseWithResult_EmptyInput(t *testing.T) {
	ctx := t.Context()
	p, _ := NewMarkdownParser(GoMarkdown)
	res := p.ParseWithResult(ctx, "empty.md", []byte(""))
	if res.Err != nil {
		t.Fatalf("ParseWithResult: %v", res.Err)
	}
	if len(res.JSON) != 1 {
		t.Fatalf("len(JSON) = %d, want 1", len(res.JSON))
	}
}

func TestMarkdownParser_ParseWithResult_ImageDataURI(t *testing.T) {
	ctx := t.Context()
	p, _ := NewMarkdownParser(GoMarkdown)
	// 1×1 pixel transparent PNG encoded as data URI
	pixelPNG := make([]byte, 68) // minimal 1x1 PNG header
	pixelB64 := base64.StdEncoding.EncodeToString([]byte("fake-png-data"))
	md := "Some text with an image\n![test](data:image/png;base64," + pixelB64 + ")\n"
	_ = pixelPNG
	res := p.ParseWithResult(ctx, "test.md", []byte(md))
	if res.Err != nil {
		t.Fatalf("ParseWithResult: %v", res.Err)
	}
	// Find the image item
	found := false
	for _, item := range res.JSON {
		if kd, _ := item["doc_type_kwd"].(string); kd == "image" {
			found = true
			if img, ok := item["image"].(string); !ok || img != pixelB64 {
				t.Fatalf("image data = %q, want %q", img, pixelB64)
			}
			break
		}
	}
	if !found {
		t.Fatal("no item with doc_type_kwd == 'image' found")
	}
}

func TestMarkdownParser_ParseWithResult_NoImage(t *testing.T) {
	ctx := t.Context()
	p, _ := NewMarkdownParser(GoMarkdown)
	md := "# Title\n\nJust some text, no images here.\n\nMore text."
	res := p.ParseWithResult(ctx, "test.md", []byte(md))
	if res.Err != nil {
		t.Fatalf("ParseWithResult: %v", res.Err)
	}
	for _, item := range res.JSON {
		if kd, _ := item["doc_type_kwd"].(string); kd == "image" {
			t.Fatal("unexpected image item in text-only markdown")
		}
	}
}

func TestMarkdownParser_ParseWithResult_RendersTableInline(t *testing.T) {
	ctx := t.Context()
	p, _ := NewMarkdownParser(GoMarkdown)
	md := "[M03] Health check package comparison:\n\n| Check item | Basic 699 CNY | Advanced 1299 CNY |\n| --- | --- | --- |\n| Blood routine / Urine routine | Yes | Yes |\n\nNote: All packages require fasting.\n"
	res := p.ParseWithResult(ctx, "test.md", []byte(md))
	if res.Err != nil {
		t.Fatalf("ParseWithResult: %v", res.Err)
	}
	if len(res.JSON) != 1 {
		t.Fatalf("len(JSON) = %d, want 1", len(res.JSON))
	}
	text, _ := res.JSON[0]["text"].(string)
	if got, _ := res.JSON[0]["doc_type_kwd"].(string); got != "text" {
		t.Fatalf("doc_type_kwd = %q, want text", got)
	}
	if !strings.Contains(text, "<table>") || !strings.Contains(text, "<th>Check item</th>") {
		t.Fatalf("table was not rendered inline: %q", text)
	}
	if strings.Contains(text, "| Check item |") {
		t.Fatalf("raw markdown table leaked into text: %q", text)
	}
	if gap := text[strings.Index(text, "</table>"):strings.Index(text, "Note:")]; gap != "</table>\n" {
		t.Fatalf("gap after table = %q, want %q", gap, "</table>\n")
	}
}

func TestMarkdownParser_ConfigureFromSetup(t *testing.T) {
	p, _ := NewMarkdownParser(GoMarkdown)
	p.ConfigureFromSetup(map[string]any{
		"parse_method":          "deepdoc",
		"output_format":         "json",
		"vlm":                   map[string]any{"llm_id": "gpt-4-vision"},
		"flatten_media_to_text": false,
	})
	if p.ParseMethod != "deepdoc" {
		t.Fatalf("ParseMethod = %q, want %q", p.ParseMethod, "deepdoc")
	}
	if p.OutputFormat != "json" {
		t.Fatalf("OutputFormat = %q, want %q", p.OutputFormat, "json")
	}
	if p.VLM == nil {
		t.Fatal("VLM is nil; want map")
	}
	if id, _ := p.VLM["llm_id"].(string); id != "gpt-4-vision" {
		t.Fatalf("VLM[llm_id] = %q, want %q", id, "gpt-4-vision")
	}
}

func TestMarkdownParser_ConfigureFromSetup_NilSafe(t *testing.T) {
	p, _ := NewMarkdownParser(GoMarkdown)
	p.ConfigureFromSetup(nil) // should not panic
	if p.ParseMethod != "" {
		t.Fatalf("ParseMethod should be empty after nil setup, got %q", p.ParseMethod)
	}
}

// TestMarkdownParser_FlattenMediaToText verifies that when
// flatten_media_to_text is true, image items are emitted with
// doc_type_kwd="text" (mirroring Python parser.py:1034). When false,
// image items keep doc_type_kwd="image".
func TestMarkdownParser_FlattenMediaToText(t *testing.T) {
	ctx := t.Context()
	pixelB64 := base64.StdEncoding.EncodeToString([]byte("fake-png-data"))
	md := "Some text with an image\n![test](data:image/png;base64," + pixelB64 + ")\n"

	cases := []struct {
		name        string
		flatten     bool
		wantDocType string
	}{
		{"flatten=false keeps image", false, "image"},
		{"flatten=true forces text", true, "text"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			p, _ := NewMarkdownParser(GoMarkdown)
			p.ConfigureFromSetup(map[string]any{"flatten_media_to_text": tc.flatten})
			if p.FlattenMediaToText != tc.flatten {
				t.Fatalf("FlattenMediaToText field = %v, want %v", p.FlattenMediaToText, tc.flatten)
			}
			res := p.ParseWithResult(ctx, "test.md", []byte(md))
			if res.Err != nil {
				t.Fatalf("ParseWithResult: %v", res.Err)
			}
			for _, item := range res.JSON {
				if kd, _ := item["doc_type_kwd"].(string); kd != tc.wantDocType {
					t.Errorf("doc_type_kwd = %q, want %q (item text=%q)",
						kd, tc.wantDocType, item["text"])
				}
			}
		})
	}
}

func TestResolveMarkdownImage_DataURI(t *testing.T) {
	b64 := base64.StdEncoding.EncodeToString([]byte("fakeimage"))
	md := "![alt](data:image/png;base64," + b64 + ")"
	result, found := resolveMarkdownImage("", md)
	if !found {
		t.Fatal("expected image found for data URI")
	}
	if result != b64 {
		t.Fatalf("got %q, want %q", result, b64)
	}
}

func TestResolveMarkdownImage_NoImage(t *testing.T) {
	_, found := resolveMarkdownImage("", "# Hello\nNo image here")
	if found {
		t.Fatal("expected no image found")
	}
}

func TestResolveMarkdownImage_HTTPImage(t *testing.T) {
	// httptest servers bind loopback, which the SSRF guard rejects by
	// default. Allow loopback for this test so the HTTP fetch path is
	// exercised (production keeps ssrfAllowLoopback == false).
	prev := ssrfAllowLoopback
	ssrfAllowLoopback = true
	defer func() { ssrfAllowLoopback = prev }()

	// Start a test HTTP server serving a fake PNG
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		w.Write([]byte("fake-png-bytes"))
	}))
	defer ts.Close()

	md := "![alt](" + ts.URL + "/image.png)"
	result, found := resolveMarkdownImage("", md)
	if !found {
		t.Fatal("expected image found for HTTP URL")
	}
	expectedB64 := base64.StdEncoding.EncodeToString([]byte("fake-png-bytes"))
	if result != expectedB64 {
		t.Fatalf("got %q, want %q", result, expectedB64)
	}
}

func TestFetchImageAsBase64_RejectsCredentials(t *testing.T) {
	_, err := fetchImageAsBase64("https://user:pass@example.com/img.png")
	if err == nil {
		t.Fatal("expected error for URL with credentials")
	}
}

func TestFetchImageAsBase64_InvalidURL(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer ts.Close()

	_, err := fetchImageAsBase64(ts.URL + "/nonexistent.png")
	if err == nil {
		t.Fatal("expected error for 404 response")
	}
}