Enterprise-Knowledge-Base/internal/handler/dto/model_test.go

package dto

import (
	"context"
	"encoding/json"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestModelResponse_OmitsSecrets(t *testing.T) {
	m := &types.Model{
		ID:          "m-1",
		Name:        "gpt-x",
		DisplayName: "Support QA",
		Parameters: types.ModelParameters{
			APIKey:    "sk-real-api-key-do-not-leak",
			AppSecret: "app-real-secret-do-not-leak",
			AppID:     "appid-public-ok-to-show",
			BaseURL:   "https://api.example.com",
			Provider:  "openai",
		},
	}
	body, err := json.Marshal(NewModelResponse(adminContext(), m))
	assert.NoError(t, err)
	s := string(body)
	assert.NotContains(t, s, "sk-real-api-key-do-not-leak")
	assert.NotContains(t, s, "app-real-secret-do-not-leak")
	// Parameters sub-object must contain no secret keys.
	var raw map[string]json.RawMessage
	assert.NoError(t, json.Unmarshal(body, &raw))
	params := string(raw["parameters"])
	assert.NotContains(t, params, `"api_key"`)
	assert.NotContains(t, params, `"app_secret"`)
	// Credential metadata map exposes booleans only.
	assert.Contains(t, s, `"credentials"`)
	assert.Contains(t, s, `"api_key":{"configured":true}`)
	assert.Contains(t, s, `"app_secret":{"configured":true}`)
	// Non-secret fields pass through.
	assert.Contains(t, s, "appid-public-ok-to-show")
	assert.Contains(t, s, "api.example.com")
	assert.Contains(t, s, `"display_name":"Support QA"`)
}

func TestModelResponse_BuiltinStripsTenantConfig(t *testing.T) {
	m := &types.Model{
		ID:        "builtin-1",
		IsBuiltin: true,
		Parameters: types.ModelParameters{
			BaseURL:        "https://tenant-private.example.com",
			APIKey:         "should-not-leak",
			AppID:          "tenant-app-id",
			SupportsVision: true,
			ExtraConfig:    map[string]string{"region": "cn-hangzhou"},
		},
	}
	resp := NewModelResponse(adminContext(), m)
	assert.Empty(t, resp.Parameters.BaseURL,
		"builtin must not leak per-tenant base URL")
	assert.Empty(t, resp.Parameters.AppID,
		"builtin must not leak per-tenant app_id")
	assert.Nil(t, resp.Parameters.ExtraConfig,
		"builtin must not leak per-tenant extra_config")
	assert.True(t, resp.Parameters.SupportsVision,
		"capability metadata must survive (not per-tenant)")

	body, _ := json.Marshal(resp)
	assert.False(t, strings.Contains(string(body), "should-not-leak"))
	assert.False(t, strings.Contains(string(body), "tenant-private.example.com"))
}

func TestModelResponse_SystemAdminCanManageBuiltinConfig(t *testing.T) {
	ctx := context.WithValue(viewerContext(), types.SystemAdminContextKey, true)
	m := &types.Model{
		ID:        "builtin-1",
		IsBuiltin: true,
		Parameters: types.ModelParameters{
			BaseURL:       "https://global-provider.example.com",
			APIKey:        "should-never-be-returned",
			AppSecret:     "also-never-returned",
			AppID:         "global-app-id",
			ExtraConfig:   map[string]string{"region": "ap-guangzhou"},
			CustomHeaders: map[string]string{"X-Route": "global"},
		},
	}

	resp := NewModelResponse(ctx, m)
	assert.Equal(t, "https://global-provider.example.com", resp.Parameters.BaseURL)
	assert.Equal(t, "global-app-id", resp.Parameters.AppID)
	assert.Equal(t, map[string]string{"region": "ap-guangzhou"}, resp.Parameters.ExtraConfig)
	assert.Equal(t, map[string]string{"X-Route": "global"}, resp.Parameters.CustomHeaders)
	assert.True(t, resp.Credentials["api_key"].Configured)
	assert.True(t, resp.Credentials["app_secret"].Configured)

	body, err := json.Marshal(resp)
	require.NoError(t, err)
	assert.NotContains(t, string(body), "should-never-be-returned")
	assert.NotContains(t, string(body), "also-never-returned")
}

func TestModelResponse_ViewerStripsIntegrationDetail(t *testing.T) {
	m := &types.Model{
		ID: "m-2",
		Parameters: types.ModelParameters{
			BaseURL:       "https://tenant-private.example.com",
			CustomHeaders: map[string]string{"Authorization": "Bearer secret"},
			ExtraConfig:   map[string]string{"region": "cn-hangzhou"},
		},
	}
	resp := NewModelResponse(viewerContext(), m)
	assert.Empty(t, resp.Parameters.BaseURL)
	assert.Nil(t, resp.Parameters.CustomHeaders)
	assert.Nil(t, resp.Parameters.ExtraConfig)
}

func TestModelResponse_NilSafe(t *testing.T) {
	assert.Nil(t, NewModelResponse(adminContext(), nil))
	assert.Equal(t, []*ModelResponse{}, NewModelResponses(adminContext(), nil))
}