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) {
	dbPath := "test_db"
	defer os.RemoveAll(dbPath)

	dbUnauth, err := GetDB(dbPath, 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")
	}

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

	// Wait and test search after rebuild
	resultsRebuilt, err := db.Search("", "火星探测", 10, nil)
	if err != nil {
		t.Fatalf("Search after rebuild failed: %v", err)
	}
	if len(resultsRebuilt) == 0 || resultsRebuilt[0].ID != "1" {
		t.Fatalf("Expected doc1 after rebuild")
	}

}