80 lines
2.2 KiB
Go
80 lines
2.2 KiB
Go
package ai
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
"senlinai-agent/backend/internal/domain"
|
|
)
|
|
|
|
func TestSelectKeyPrefersUserKey(t *testing.T) {
|
|
database := newTestDB(t)
|
|
require.NoError(t, database.Create(&domain.AIKey{UserID: 3, Provider: "openai", EncryptedAPIKey: "user-key"}).Error)
|
|
gateway := NewGateway(database, "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 TestSelectKeyFallsBackToSystemKey(t *testing.T) {
|
|
database := newTestDB(t)
|
|
gateway := NewGateway(database, "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) {
|
|
database := newTestDB(t)
|
|
gateway := NewGateway(database, "")
|
|
|
|
_, err := gateway.SelectKey(3)
|
|
|
|
require.ErrorContains(t, err, "no ai key available")
|
|
}
|
|
|
|
func TestRecordCallStoresAuditFields(t *testing.T) {
|
|
database := newTestDB(t)
|
|
gateway := NewGateway(database, "system-key")
|
|
|
|
require.NoError(t, gateway.RecordCall(3, "openai", "system", "inbox_analyze", "failed", "rate limited"))
|
|
|
|
var log domain.AICallLog
|
|
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 TestCreateAISession(t *testing.T) {
|
|
database := newTestDB(t)
|
|
service := NewSessionService(database)
|
|
|
|
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, domain.AutoMigrate(database))
|
|
return database
|
|
}
|