mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-12 06:35:44 +08:00
462 lines
13 KiB
Go
462 lines
13 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package task
|
|
|
|
import (
|
|
"encoding/base64"
|
|
"errors"
|
|
"regexp"
|
|
"testing"
|
|
)
|
|
|
|
const (
|
|
testDocID = "doc-1"
|
|
testKBID = "kb-1"
|
|
testDocName = "test.pdf"
|
|
)
|
|
|
|
func TestPrepareDocs_Basic(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Hello World", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
if len(chunks) != 1 {
|
|
t.Fatalf("expected 1 chunk, got %d", len(chunks))
|
|
}
|
|
|
|
c := chunks[0]
|
|
if c.ContentWithWeight != "Hello World" {
|
|
t.Errorf("ContentWithWeight = %q, want %q", c.ContentWithWeight, "Hello World")
|
|
}
|
|
if c.DocID != testDocID {
|
|
t.Errorf("DocID = %q, want %q", c.DocID, testDocID)
|
|
}
|
|
if c.KBID != testKBID {
|
|
t.Errorf("KBID = %q, want %q", c.KBID, testKBID)
|
|
}
|
|
if c.DocNameKwd != testDocName {
|
|
t.Errorf("DocNameKwd = %q, want %q", c.DocNameKwd, testDocName)
|
|
}
|
|
if c.AvailableInt != 1 {
|
|
t.Errorf("AvailableInt = %d, want 1", c.AvailableInt)
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_Multiple(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "First", "doc_type_kwd": "text", "img_id": ""},
|
|
{"text": "Second", "doc_type_kwd": "text", "img_id": ""},
|
|
{"text": "Third", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
if len(chunks) != 3 {
|
|
t.Fatalf("expected 3 chunks, got %d", len(chunks))
|
|
}
|
|
|
|
// All IDs should be different (different content)
|
|
ids := make(map[string]bool)
|
|
for _, c := range chunks {
|
|
if c.ContentWithWeight == "" {
|
|
t.Error("chunk has empty ContentWithWeight")
|
|
}
|
|
ids[c.ID] = true
|
|
}
|
|
if len(ids) != 3 {
|
|
t.Errorf("expected 3 unique IDs, got %d", len(ids))
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_EmptyInput(t *testing.T) {
|
|
chunks := PrepareDocs(nil, testDocID, testKBID, testDocName, 0)
|
|
if chunks != nil {
|
|
t.Errorf("expected nil for empty input, got %d chunks", len(chunks))
|
|
}
|
|
|
|
chunks = PrepareDocs([]map[string]any{}, testDocID, testKBID, testDocName, 0)
|
|
if chunks != nil {
|
|
t.Errorf("expected nil for empty input, got %d chunks", len(chunks))
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_IDFormat(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Content for ID test", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
id := chunks[0].ID
|
|
// xxhash64 hex digest is 16 hex chars (64 bits = 16 hex digits)
|
|
matched, _ := regexp.MatchString(`^[0-9a-f]{16}$`, id)
|
|
if !matched {
|
|
t.Errorf("ID %q does not match hex format", id)
|
|
}
|
|
|
|
}
|
|
|
|
func TestPrepareDocs_IDDeterministic(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Deterministic test content", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
|
|
chunks1 := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
chunks2 := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
if chunks1[0].ID != chunks2[0].ID {
|
|
t.Errorf("ID not deterministic: %q vs %q", chunks1[0].ID, chunks2[0].ID)
|
|
}
|
|
|
|
// Different docID should produce different ID
|
|
chunks3 := PrepareDocs(sections, "doc-2", testKBID, testDocName, 0)
|
|
if chunks1[0].ID == chunks3[0].ID {
|
|
t.Error("different docIDs produced the same chunk ID")
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_Timestamps(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Timestamp test", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
c := chunks[0]
|
|
|
|
// CreateTime should match the expected format
|
|
matched, _ := regexp.MatchString(`^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}$`, c.CreateTime)
|
|
if !matched {
|
|
t.Errorf("CreateTime %q does not match format YYYY-MM-DD HH:MM:SS", c.CreateTime)
|
|
}
|
|
|
|
if c.CreateTimestamp <= 0 {
|
|
t.Errorf("CreateTimestamp = %f, want > 0", c.CreateTimestamp)
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_Positions(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{
|
|
"text": "Positioned content",
|
|
"doc_type_kwd": "text",
|
|
"img_id": "",
|
|
"positions": []float64{3, 150.5, 20, 400, 60},
|
|
},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
c := chunks[0]
|
|
if len(c.PageNumInt) != 1 || c.PageNumInt[0] != 3 {
|
|
t.Errorf("PageNumInt = %v, want [3]", c.PageNumInt)
|
|
}
|
|
if len(c.TopInt) != 1 || c.TopInt[0] != 150 {
|
|
t.Errorf("TopInt = %v, want [150]", c.TopInt)
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_NoPositions(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "No positions", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
c := chunks[0]
|
|
if c.PageNumInt != nil {
|
|
t.Errorf("PageNumInt = %v, want nil", c.PageNumInt)
|
|
}
|
|
if c.TopInt != nil {
|
|
t.Errorf("TopInt = %v, want nil", c.TopInt)
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_ImgID(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{
|
|
"text": "Content with image",
|
|
"doc_type_kwd": "text",
|
|
"img_id": "some_image_id_123",
|
|
},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
c := chunks[0]
|
|
if c.ImgID != "some_image_id_123" {
|
|
t.Errorf("ImgID = %q, want %q", c.ImgID, "some_image_id_123")
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_EmptyImgID(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "No image", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
c := chunks[0]
|
|
if c.ImgID != "" {
|
|
t.Errorf("ImgID = %q, want empty string", c.ImgID)
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_MixedContent(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Page 1 content", "doc_type_kwd": "text", "img_id": "",
|
|
"positions": []float64{1, 0, 0, 100, 50}},
|
|
{"text": "Page 2 with image", "doc_type_kwd": "text", "img_id": "img_abc",
|
|
"positions": []float64{2, 30, 10, 200, 80}},
|
|
{"text": "No position, no image", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
if len(chunks) != 3 {
|
|
t.Fatalf("expected 3 chunks, got %d", len(chunks))
|
|
}
|
|
|
|
// First: has position
|
|
if len(chunks[0].PageNumInt) != 1 || chunks[0].PageNumInt[0] != 1 {
|
|
t.Errorf("chunk[0] PageNumInt = %v, want [1]", chunks[0].PageNumInt)
|
|
}
|
|
|
|
// Second: has position + img_id
|
|
if chunks[1].ImgID != "img_abc" {
|
|
t.Errorf("chunk[1] ImgID = %q, want %q", chunks[1].ImgID, "img_abc")
|
|
}
|
|
if len(chunks[1].PageNumInt) != 1 || chunks[1].PageNumInt[0] != 2 {
|
|
t.Errorf("chunk[1] PageNumInt = %v, want [2]", chunks[1].PageNumInt)
|
|
}
|
|
|
|
// Third: no position, no image
|
|
if chunks[2].PageNumInt != nil {
|
|
t.Errorf("chunk[2] PageNumInt should be nil, got %v", chunks[2].PageNumInt)
|
|
}
|
|
if chunks[2].ImgID != "" {
|
|
t.Errorf("chunk[2] ImgID should be empty, got %q", chunks[2].ImgID)
|
|
}
|
|
}
|
|
|
|
func TestPrepareDocs_PageRank(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Ranked content", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
|
|
// pagerank=0 → PageRank should be 0 (unset)
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
if chunks[0].PageRank != 0 {
|
|
t.Errorf("PageRank = %d, want 0", chunks[0].PageRank)
|
|
}
|
|
|
|
// pagerank=5 → PageRank should be 5
|
|
chunks = PrepareDocs(sections, testDocID, testKBID, testDocName, 5)
|
|
if chunks[0].PageRank != 5 {
|
|
t.Errorf("PageRank = %d, want 5", chunks[0].PageRank)
|
|
}
|
|
}
|
|
|
|
func TestUploadChunkImages_WithImage(t *testing.T) {
|
|
imgBase64 := base64.StdEncoding.EncodeToString([]byte("fake image"))
|
|
sections := []map[string]any{
|
|
{"text": "Has image", "doc_type_kwd": "text", "img_id": "", "image": imgBase64},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
var putBucket, putFnm string
|
|
putFn := func(bucket, fnm string, binary []byte) error {
|
|
putBucket = bucket
|
|
putFnm = fnm
|
|
if len(binary) == 0 {
|
|
t.Error("Put called with empty binary")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil {
|
|
t.Fatalf("UploadChunkImages failed: %v", err)
|
|
}
|
|
|
|
expectedImgID := testKBID + "-" + chunks[0].ID
|
|
if chunks[0].ImgID != expectedImgID {
|
|
t.Errorf("ImgID = %q, want %q", chunks[0].ImgID, expectedImgID)
|
|
}
|
|
if putBucket != testKBID {
|
|
t.Errorf("put bucket = %q, want %q", putBucket, testKBID)
|
|
}
|
|
if putFnm != chunks[0].ID {
|
|
t.Errorf("put fnm = %q, want %q", putFnm, chunks[0].ID)
|
|
}
|
|
}
|
|
|
|
func TestUploadChunkImages_PresetImgID(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Has preset img_id", "doc_type_kwd": "text", "img_id": "preset-id-123", "image": "should-be-ignored"},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
callCount := 0
|
|
putFn := func(bucket, fnm string, binary []byte) error {
|
|
callCount++
|
|
return nil
|
|
}
|
|
|
|
if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil {
|
|
t.Fatalf("UploadChunkImages failed: %v", err)
|
|
}
|
|
|
|
if chunks[0].ImgID != "preset-id-123" {
|
|
t.Errorf("ImgID = %q, want %q", chunks[0].ImgID, "preset-id-123")
|
|
}
|
|
if callCount != 0 {
|
|
t.Errorf("expected 0 Put calls (preset img_id), got %d", callCount)
|
|
}
|
|
}
|
|
|
|
func TestUploadChunkImages_NoImage(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "No image", "doc_type_kwd": "text", "img_id": ""},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
callCount := 0
|
|
putFn := func(bucket, fnm string, binary []byte) error {
|
|
callCount++
|
|
return nil
|
|
}
|
|
|
|
if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil {
|
|
t.Fatalf("UploadChunkImages failed: %v", err)
|
|
}
|
|
|
|
if chunks[0].ImgID != "" {
|
|
t.Errorf("ImgID = %q, want empty string", chunks[0].ImgID)
|
|
}
|
|
if callCount != 0 {
|
|
t.Errorf("expected 0 Put calls, got %d", callCount)
|
|
}
|
|
}
|
|
|
|
func TestUploadChunkImages_Mixed(t *testing.T) {
|
|
imgBase64 := base64.StdEncoding.EncodeToString([]byte("img1"))
|
|
sections := []map[string]any{
|
|
{"text": "Img1", "doc_type_kwd": "text", "img_id": "", "image": imgBase64},
|
|
{"text": "No img", "doc_type_kwd": "text", "img_id": ""},
|
|
{"text": "Preset", "doc_type_kwd": "text", "img_id": "preset-001"},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
var callCount int
|
|
putFn := func(bucket, fnm string, binary []byte) error {
|
|
callCount++
|
|
return nil
|
|
}
|
|
|
|
if err := UploadChunkImages(sections, chunks, testKBID, putFn); err != nil {
|
|
t.Fatalf("UploadChunkImages failed: %v", err)
|
|
}
|
|
|
|
if callCount != 1 {
|
|
t.Errorf("expected 1 Put call, got %d", callCount)
|
|
}
|
|
if chunks[0].ImgID != testKBID+"-"+chunks[0].ID {
|
|
t.Errorf("chunk[0] ImgID = %q, want %q", chunks[0].ImgID, testKBID+"-"+chunks[0].ID)
|
|
}
|
|
if chunks[1].ImgID != "" {
|
|
t.Errorf("chunk[1] ImgID = %q, want empty", chunks[1].ImgID)
|
|
}
|
|
if chunks[2].ImgID != "preset-001" {
|
|
t.Errorf("chunk[2] ImgID = %q, want 'preset-001'", chunks[2].ImgID)
|
|
}
|
|
}
|
|
|
|
func TestUploadChunkImages_InvalidBase64(t *testing.T) {
|
|
sections := []map[string]any{
|
|
{"text": "Bad image", "doc_type_kwd": "text", "img_id": "", "image": "!!!not-base64!!!"},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
putFn := func(bucket, fnm string, binary []byte) error {
|
|
return nil
|
|
}
|
|
|
|
if err := UploadChunkImages(sections, chunks, testKBID, putFn); err == nil {
|
|
t.Fatal("expected error for invalid base64, got nil")
|
|
}
|
|
}
|
|
|
|
func TestUploadChunkImages_PutError(t *testing.T) {
|
|
imgBase64 := base64.StdEncoding.EncodeToString([]byte("img"))
|
|
sections := []map[string]any{
|
|
{"text": "Put fails", "doc_type_kwd": "text", "img_id": "", "image": imgBase64},
|
|
}
|
|
chunks := PrepareDocs(sections, testDocID, testKBID, testDocName, 0)
|
|
|
|
putFn := func(bucket, fnm string, binary []byte) error {
|
|
return errors.New("storage unavailable")
|
|
}
|
|
|
|
if err := UploadChunkImages(sections, chunks, testKBID, putFn); err == nil {
|
|
t.Fatal("expected error for put failure, got nil")
|
|
}
|
|
}
|
|
|
|
func TestChunkID_NotPanic(t *testing.T) {
|
|
// Edge cases: empty content
|
|
_ = ChunkID("", testDocID)
|
|
_ = ChunkID("content", "")
|
|
_ = ChunkID("", "")
|
|
// All should not panic
|
|
}
|
|
|
|
func TestChunkID_PreservesLeadingZero(t *testing.T) {
|
|
content := ""
|
|
docID := "doc-1"
|
|
got := ChunkID(content, docID)
|
|
want := "037fe13bd80c56aa"
|
|
if got != want {
|
|
t.Fatalf("ChunkID(%q, %q) = %q, want %q", content, docID, got, want)
|
|
}
|
|
if len(got) != 16 {
|
|
t.Fatalf("ChunkID length = %d, want 16", len(got))
|
|
}
|
|
}
|
|
|
|
func TestChunk_String(t *testing.T) {
|
|
c := Chunk{
|
|
ID: "abc123",
|
|
DocID: "doc-1",
|
|
ContentWithWeight: "Short text",
|
|
PageNumInt: []int{1},
|
|
TopInt: []int{0},
|
|
}
|
|
s := c.String()
|
|
if len(s) == 0 {
|
|
t.Error("String() returned empty")
|
|
}
|
|
}
|
|
|
|
func TestChunk_StringTruncates(t *testing.T) {
|
|
long := ""
|
|
for i := 0; i < 200; i++ {
|
|
long += "x"
|
|
}
|
|
c := Chunk{
|
|
ID: "abc123",
|
|
DocID: "doc-1",
|
|
ContentWithWeight: long,
|
|
}
|
|
s := c.String()
|
|
if len(s) >= 200 {
|
|
t.Error("String() should truncate long content")
|
|
}
|
|
}
|