107 lines
3.0 KiB
Go
107 lines
3.0 KiB
Go
package ai
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
"senlinai-agent/backend/internal/models"
|
|
)
|
|
|
|
func TestSelectKeyPrefersUserKey(t *testing.T) {
|
|
database := newTestDB(t)
|
|
require.NoError(t, database.Create(&models.SenlinAgentAIKey{UserID: 3, Provider: "openai", EncryptedAPIKey: "user-key"}).Error)
|
|
gateway := NewGateway("system-key")
|
|
|
|
selected, err := gateway.SelectKey(3)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, "user", selected.KeyType)
|
|
require.Equal(t, "user-key", selected.APIKey)
|
|
}
|
|
|
|
func TestSaveUserKeyEncryptsStoredKey(t *testing.T) {
|
|
database := newTestDB(t)
|
|
gateway := NewGatewayWithSecret("system-key", "test-encryption-secret")
|
|
|
|
require.NoError(t, gateway.SaveUserKey(3, "openai", "user-key"))
|
|
|
|
var stored models.SenlinAgentAIKey
|
|
require.NoError(t, database.Where("user_id = ?", 3).First(&stored).Error)
|
|
require.NotEqual(t, "user-key", stored.EncryptedAPIKey)
|
|
require.Contains(t, stored.EncryptedAPIKey, "v1:")
|
|
selected, err := gateway.SelectKey(3)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "user-key", selected.APIKey)
|
|
}
|
|
|
|
func TestSelectKeyFallsBackToSystemKey(t *testing.T) {
|
|
newTestDB(t)
|
|
gateway := NewGateway("system-key")
|
|
|
|
selected, err := gateway.SelectKey(3)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, "system", selected.KeyType)
|
|
require.Equal(t, "system-key", selected.APIKey)
|
|
}
|
|
|
|
func TestSelectKeyReturnsErrorWhenNoKeyAvailable(t *testing.T) {
|
|
newTestDB(t)
|
|
gateway := NewGateway("")
|
|
|
|
_, err := gateway.SelectKey(3)
|
|
|
|
require.ErrorContains(t, err, "no ai key available")
|
|
}
|
|
|
|
func TestRecordCallStoresAuditFields(t *testing.T) {
|
|
database := newTestDB(t)
|
|
gateway := NewGateway("system-key")
|
|
|
|
require.NoError(t, gateway.RecordCall(3, "openai", "system", "inbox_analyze", "failed", "rate limited"))
|
|
|
|
var log models.SenlinAgentAICallLog
|
|
require.NoError(t, database.First(&log).Error)
|
|
require.Equal(t, uint(3), log.UserID)
|
|
require.Equal(t, "openai", log.Provider)
|
|
require.Equal(t, "system", log.UsedKeyType)
|
|
require.Equal(t, "inbox_analyze", log.Action)
|
|
require.Equal(t, "failed", log.Status)
|
|
require.Equal(t, "rate limited", log.Error)
|
|
}
|
|
|
|
func TestCheckRateLimitRejectsCallsOverWindow(t *testing.T) {
|
|
newTestDB(t)
|
|
gateway := NewGateway("system-key")
|
|
require.NoError(t, gateway.RecordCall(3, "openai", "system", "inbox_analyze", "succeeded", ""))
|
|
|
|
err := gateway.CheckRateLimit(3, "inbox_analyze", 1, time.Hour)
|
|
|
|
require.ErrorContains(t, err, "ai rate limit exceeded")
|
|
}
|
|
|
|
func TestCreateAISession(t *testing.T) {
|
|
newTestDB(t)
|
|
service := NewSessionService()
|
|
|
|
session, err := service.Create(7, 3, "报价分析")
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, uint(7), session.ProjectID)
|
|
require.Equal(t, uint(3), session.CreatedBy)
|
|
require.Equal(t, "报价分析", session.Title)
|
|
}
|
|
|
|
func newTestDB(t *testing.T) *gorm.DB {
|
|
t.Helper()
|
|
database, err := gorm.Open(sqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
require.NoError(t, models.AutoMigrate(database))
|
|
models.DBService = database
|
|
return database
|
|
}
|