ragflow/internal/handler/dataset_list_test.go

package handler

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"gorm.io/gorm"

	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/entity"
	dataset "ragflow/internal/service/dataset"
)

const (
	listDatasetsTestKBID       = "123e4567e89b12d3a456426614174000"
	listDatasetsTestKBIDDashed = "123e4567-e89b-12d3-a456-426614174000"
)

func setupListDatasetsTestDB(t *testing.T) *gorm.DB {
	t.Helper()

	db, err := gorm.Open(sqlite.Open("file:"+url.QueryEscape(t.Name())+"?mode=memory&cache=shared"), &gorm.Config{
		TranslateError: true,
	})
	if err != nil {
		t.Fatalf("failed to open sqlite: %v", err)
	}

	if err := db.AutoMigrate(&entity.Knowledgebase{}, &entity.User{}, &entity.UserTenant{}, &entity.Tenant{}); err != nil {
		t.Fatalf("failed to migrate test schema: %v", err)
	}

	origDB := dao.DB
	dao.DB = db
	t.Cleanup(func() { dao.DB = origDB })

	return db
}

func insertListDatasetsTestKB(t *testing.T, id, tenantID, name string) {
	t.Helper()

	status := string(entity.StatusValid)
	kb := &entity.Knowledgebase{
		ID:           id,
		TenantID:     tenantID,
		Name:         name,
		EmbdID:       "BAAI/bge-large-zh-v1.5@Builtin",
		CreatedBy:    tenantID,
		Permission:   string(entity.TenantPermissionMe),
		ParserID:     string(entity.ParserTypeNaive),
		ParserConfig: entity.JSONMap{"chunk_token_num": float64(128)},
		Status:       &status,
	}
	if err := dao.DB.Create(kb).Error; err != nil {
		t.Fatalf("insert test kb: %v", err)
	}
}

func newListDatasetsTestRouter() *gin.Engine {
	gin.SetMode(gin.TestMode)
	h := NewDatasetsHandler(dataset.NewDatasetService(), nil)
	r := gin.New()
	r.GET("/api/v1/datasets", func(c *gin.Context) {
		c.Set("user", &entity.User{ID: "user-1"})
		h.ListDatasets(c)
	})
	return r
}

type listDatasetsTestResponse struct {
	Code          int                      `json:"code"`
	Data          []map[string]interface{} `json:"data"`
	Message       string                   `json:"message"`
	TotalDatasets int64                    `json:"total_datasets"`
}

func getListDatasets(t *testing.T, r *gin.Engine, rawQuery string) listDatasetsTestResponse {
	t.Helper()

	resp := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/api/v1/datasets?"+rawQuery, nil)
	r.ServeHTTP(resp, req)

	if resp.Code != http.StatusOK {
		t.Fatalf("status=%d body=%s", resp.Code, resp.Body.String())
	}

	var body listDatasetsTestResponse
	if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil {
		t.Fatalf("unmarshal response: %v body=%s", err, resp.Body.String())
	}
	return body
}

func TestDatasetsHandlerListDatasetsFiltersByIDs(t *testing.T) {
	setupListDatasetsTestDB(t)
	insertListDatasetsTestKB(t, listDatasetsTestKBID, "user-1", "Alpha")

	body := getListDatasets(t, newListDatasetsTestRouter(),
		fmt.Sprintf("ids=%s&page_size=1", listDatasetsTestKBIDDashed))

	if body.Code != int(common.CodeSuccess) {
		t.Fatalf("code=%d message=%q", body.Code, body.Message)
	}
	if body.TotalDatasets != 1 || len(body.Data) != 1 {
		t.Fatalf("expected exactly one dataset, got total=%d len=%d", body.TotalDatasets, len(body.Data))
	}
	if body.Data[0]["id"] != listDatasetsTestKBID {
		t.Fatalf("expected dataset id %q, got %#v", listDatasetsTestKBID, body.Data[0]["id"])
	}
}

func TestDatasetsHandlerListDatasetsRejectsInvalidIDInIDs(t *testing.T) {
	setupListDatasetsTestDB(t)

	body := getListDatasets(t, newListDatasetsTestRouter(), "ids=not-a-uuid")

	if body.Code != int(common.CodeArgumentError) {
		t.Fatalf("code=%d want=%d", body.Code, common.CodeArgumentError)
	}
	if body.Message != "Invalid UUID format" {
		t.Fatalf("message=%q want=%q", body.Message, "Invalid UUID format")
	}
}

func TestDatasetsHandlerListDatasetsRejectsDuplicateIDs(t *testing.T) {
	setupListDatasetsTestDB(t)

	rawQuery := fmt.Sprintf("ids=%s,%s", listDatasetsTestKBIDDashed, listDatasetsTestKBIDDashed)
	body := getListDatasets(t, newListDatasetsTestRouter(), rawQuery)

	if body.Code != int(common.CodeArgumentError) {
		t.Fatalf("code=%d want=%d", body.Code, common.CodeArgumentError)
	}
	expected := fmt.Sprintf("Duplicate ids: '%s'", listDatasetsTestKBID)
	if body.Message != expected {
		t.Fatalf("message=%q want=%q", body.Message, expected)
	}
}