indexDB/indexDB_test.go

package indexDB

import (
	"os"
	"testing"
)

func mockEmbedding(text string) ([]float32, error) {
	if text == "doc1" {
		return []float32{0.1, 0.2, 0.3}, nil
	} else if text == "doc2" {
		return []float32{0.9, 0.8, 0.7}, nil
	}
	return []float32{0.0, 0.0, 0.0}, nil
}

func TestIndexDB(t *testing.T) {
	fPath := "test_fulltext"
	vPath := "test_vector"
	defer os.RemoveAll(fPath)
	defer os.RemoveAll(vPath)

	dbUnauth, err := GetDB(fPath, vPath, mockEmbedding, nil)
	if err != nil {
		t.Fatalf("Failed to create engine: %v", err)
	}

	db := dbUnauth.Auth("_system")

	err = db.Add("1", "中国航天局发射了火星探测器 doc1", map[string]any{"source": "test"}, []string{"user1"})
	if err != nil {
		t.Fatalf("Failed to add document 1: %v", err)
	}

	err = db.Add("2", "The quick brown fox jumps over the lazy dog. doc2", map[string]any{"source": "test2"}, nil)
	if err != nil {
		t.Fatalf("Failed to add document 2: %v", err)
	}

	results, err := db.Search("", "火星探测", 10, nil)
	if err != nil {
		t.Fatalf("Search failed: %v", err)
	}

	if len(results) == 0 {
		t.Fatalf("Expected results, got 0")
	}

	found := false
	for _, r := range results {
		if r.ID == "1" {
			found = true
			break
		}
	}
	if !found {
		t.Fatalf("Expected doc1 in results")
	}

	// Test user permissions
	user1Db := dbUnauth.Auth("user1")
	user1Results, err := user1Db.Search("", "火星探测", 10, nil)
	if err != nil {
		t.Fatalf("Search with user1 failed: %v", err)
	}
	if len(user1Results) == 0 || user1Results[0].ID != "1" {
		t.Fatalf("User1 should see doc1")
	}

	user2Db := dbUnauth.Auth("user2")
	user2Results, err := user2Db.Search("", "火星探测", 10, nil)
	if err != nil {
		t.Fatalf("Search with user2 failed: %v", err)
	}
	if len(user2Results) != 0 {
		t.Fatalf("User2 should NOT see doc1")
	}

	// Scan Test
	docs, err := db.ScanDocuments("", 10)
	if err != nil {
		t.Fatalf("Scan failed: %v", err)
	}
	if len(docs) != 2 {
		t.Fatalf("Expected 2 docs in scan, got %d", len(docs))
	}

	docsPart, err := db.ScanDocuments("1", 10)
	if err != nil {
		t.Fatalf("Scan from 1 failed: %v", err)
	}
	if len(docsPart) != 1 || docsPart[0].ID != "2" {
		t.Fatalf("Expected doc 2 in scan from 1, got %v", docsPart)
	}

	// Remove Test
	err = db.Remove("1")
	if err != nil {
		t.Fatalf("Remove failed: %v", err)
	}

	results, err = db.Search("", "火星探测", 10, nil)
	if err != nil {
		t.Fatalf("Search after remove failed: %v", err)
	}
	for _, r := range results {
		if r.ID == "1" {
			t.Fatalf("doc1 should have been removed")
		}
	}

	// Scan after remove
	docs, err = db.ScanDocuments("", 10)
	if err != nil {
		t.Fatalf("Scan after remove failed: %v", err)
	}
	if len(docs) != 1 || docs[0].ID != "2" {
		t.Fatalf("Expected 1 doc (ID 2) after remove, got %d", len(docs))
	}
}