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

package handler

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/alicebob/miniredis/v2"
	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"

	"ragflow/internal/agent/runtime"
)

func init() {
	gin.SetMode(gin.TestMode)
}

// newAdminRuntimeTestRig wires a Selector backed by miniredis and returns
// a fully-mounted gin engine with the route registered, so tests can issue
// real HTTP requests against it.
func newAdminRuntimeTestRig(t *testing.T) (*gin.Engine, *runtime.Selector, *miniredis.Miniredis) {
	t.Helper()
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	t.Cleanup(func() { _ = rdb.Close() })

	selector := runtime.NewSelector(rdb, nil)
	h := NewAdminRuntimeHandler(selector)

	eng := gin.New()
	g := eng.Group("/api/v1/admin")
	g.POST("/canvas-runtime/:tenant_id", h.SetTenantRuntime)
	return eng, selector, mr
}

func TestAdminRuntime_SetGo(t *testing.T) {
	eng, selector, _ := newAdminRuntimeTestRig(t)

	body, _ := json.Marshal(map[string]string{"runtime": "go"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/canvas-runtime/tenant_123", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	eng.ServeHTTP(w, req)

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

	var resp setRuntimeResponse
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if resp.Code != 0 || resp.TenantID != "tenant_123" || resp.Runtime != "go" {
		t.Errorf("unexpected response: %+v", resp)
	}

	// Round-trip: the selector should now report the override.
	mode, err := selector.Select(req.Context(), "tenant_123")
	if err != nil {
		t.Fatalf("Select(): %v", err)
	}
	if mode != runtime.RuntimeGo {
		t.Errorf("Select() after SetGo = %q, want %q", mode, runtime.RuntimeGo)
	}
}

func TestAdminRuntime_SetPython(t *testing.T) {
	eng, selector, _ := newAdminRuntimeTestRig(t)

	body, _ := json.Marshal(map[string]string{"runtime": "python"})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/canvas-runtime/tenant_xyz", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	eng.ServeHTTP(w, req)

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

	mode, err := selector.Select(req.Context(), "tenant_xyz")
	if err != nil {
		t.Fatalf("Select(): %v", err)
	}
	if mode != runtime.RuntimePython {
		t.Errorf("Select() after SetPython = %q, want %q", mode, runtime.RuntimePython)
	}
}

func TestAdminRuntime_BadRequest(t *testing.T) {
	cases := []struct {
		name string
		body string
	}{
		{"unknown_mode", `{"runtime":"rust"}`},
		{"empty_mode", `{"runtime":""}`},
		{"malformed_json", `{"runtime":`},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			eng, _, _ := newAdminRuntimeTestRig(t)
			req := httptest.NewRequest(http.MethodPost,
				"/api/v1/admin/canvas-runtime/tenant_1",
				bytes.NewReader([]byte(tc.body)))
			req.Header.Set("Content-Type", "application/json")
			w := httptest.NewRecorder()
			eng.ServeHTTP(w, req)

			if w.Code != http.StatusOK {
				t.Fatalf("status = %d, want 200 envelope", w.Code)
			}
			var env map[string]any
			if err := json.Unmarshal(w.Body.Bytes(), &env); err != nil {
				t.Fatalf("decode: %v", err)
			}
			// 101 == CodeArgumentError, the only acceptable error for bad input.
			if code, _ := env["code"].(float64); code != 101 {
				t.Errorf("code = %v, want 101 (CodeArgumentError); body=%s", env["code"], w.Body.String())
			}
		})
	}
}