ragflow/internal/handler/compilation_status_test.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 handler

import (
	"encoding/json"
	"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"
)

// setupCompilationStatusHandlerDB migrates the minimal schema for the
// GET /datasets/:id/compilation/status handler and pushes it onto dao.DB.
func setupCompilationStatusHandlerDB(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.KnowledgeCompileDataset{},
	); 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 insertCompilationStatusHandlerKB(t *testing.T, kbID, ownerID string) {
	t.Helper()
	status := string(entity.StatusValid)
	kb := &entity.Knowledgebase{
		ID:         kbID,
		TenantID:   ownerID,
		Name:       "compile-status-handler-kb",
		EmbdID:     "BAAI/bge-large-zh-v1.5@Builtin",
		CreatedBy:  ownerID,
		Permission: string(entity.TenantPermissionMe),
		Status:     &status,
	}
	if err := dao.DB.Create(kb).Error; err != nil {
		t.Fatalf("insert kb: %v", err)
	}
}

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

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

func getCompilationStatus(t *testing.T, r *gin.Engine, datasetID string) (int, compilationStatusResponse) {
	t.Helper()
	resp := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet,
		"/api/v1/datasets/"+datasetID+"/compilation/status", nil)
	r.ServeHTTP(resp, req)
	var body compilationStatusResponse
	if err := json.Unmarshal(resp.Body.Bytes(), &body); err != nil {
		t.Fatalf("unmarshal response: %v body=%s", err, resp.Body.String())
	}
	return resp.Code, body
}

// TestCompilationStatusHandler_NoRowIdle verifies a dataset with no scheduling
// row returns idle with zero counts.
func TestCompilationStatusHandler_NoRowIdle(t *testing.T) {
	db := setupCompilationStatusHandlerDB(t)
	insertCompilationStatusHandlerKB(t, "kb-status-idle", "user-1")
	_ = db

	status, body := getCompilationStatus(t, newCompilationStatusHandlerRouter(), "kb-status-idle")
	if status != http.StatusOK {
		t.Fatalf("status=%d want 200", status)
	}
	if body.Code != int(common.CodeSuccess) {
		t.Fatalf("code=%d message=%q", body.Code, body.Message)
	}
	if body.Data["state"] != entity.DatasetStateIdle {
		t.Fatalf("state=%v want idle", body.Data["state"])
	}
	if n, _ := body.Data["inflight"].(float64); n != 0 {
		t.Fatalf("inflight=%v want 0", body.Data["inflight"])
	}
	if n, _ := body.Data["backlog"].(float64); n != 0 {
		t.Fatalf("backlog=%v want 0", body.Data["backlog"])
	}
}

// TestCompilationStatusHandler_FullOutput locks the JSON contract for a row
// with state, inflight/backlog counts, and error diagnostic.
func TestCompilationStatusHandler_FullOutput(t *testing.T) {
	db := setupCompilationStatusHandlerDB(t)
	insertCompilationStatusHandlerKB(t, "kb-status-full", "user-1")

	row := entity.KnowledgeCompileDataset{
		DatasetID:      "kb-status-full",
		TenantID:       "user-1",
		BacklogDocIDs:  `[{"doc_id":"d2","event_type":"completed","seq":2}]`,
		InflightDocIDs: `[{"doc_id":"d1","event_type":"completed","seq":1}]`,
		State:          entity.DatasetStatePending,
		ErrorMsg:       "merge failed: boom",
	}
	if err := db.Create(&row).Error; err != nil {
		t.Fatalf("insert scheduling row: %v", err)
	}

	status, body := getCompilationStatus(t, newCompilationStatusHandlerRouter(), "kb-status-full")
	if status != http.StatusOK {
		t.Fatalf("status=%d want 200", status)
	}
	if body.Code != int(common.CodeSuccess) {
		t.Fatalf("code=%d message=%q", body.Code, body.Message)
	}
	if body.Data["state"] != entity.DatasetStatePending {
		t.Fatalf("state=%v want pending", body.Data["state"])
	}
	if n, _ := body.Data["inflight"].(float64); n != 1 {
		t.Fatalf("inflight=%v want 1", body.Data["inflight"])
	}
	if n, _ := body.Data["backlog"].(float64); n != 1 {
		t.Fatalf("backlog=%v want 1", body.Data["backlog"])
	}
	if body.Data["error"] != "merge failed: boom" {
		t.Fatalf("error=%v want %q", body.Data["error"], "merge failed: boom")
	}
}

// TestCompilationStatusHandler_Unauthorized verifies a user who does not own the
// dataset is rejected with a data error (HTTP 200 + non-zero code, matching the
// handler's ErrorWithCode contract).
func TestCompilationStatusHandler_Unauthorized(t *testing.T) {
	db := setupCompilationStatusHandlerDB(t)
	// KB is owned by user-1; the router sets user to user-1, so this test must
	// exercise the case where the KB belongs to a different owner. We insert the
	// KB under a different owner tenant than the request user by re-pointing the
	// KB owner to "other-owner".
	insertCompilationStatusHandlerKB(t, "kb-status-forbidden", "other-owner")
	if err := db.Create(&entity.KnowledgeCompileDataset{
		DatasetID:      "kb-status-forbidden",
		TenantID:       "other-owner",
		BacklogDocIDs:  "[]",
		InflightDocIDs: "[]",
		State:          entity.DatasetStateRunning,
	}).Error; err != nil {
		t.Fatalf("insert scheduling row: %v", err)
	}

	_, body := getCompilationStatus(t, newCompilationStatusHandlerRouter(), "kb-status-forbidden")
	if body.Code != int(common.CodeDataError) {
		t.Fatalf("code=%d want %d", body.Code, common.CodeDataError)
	}
	if body.Message != "no authorization" {
		t.Fatalf("message=%q want %q", body.Message, "no authorization")
	}
}