* fix(ops): 修复运维监控系统的关键安全和稳定性问题
## 修复内容
### P0 严重问题
1. **DNS Rebinding防护** (ops_alert_service.go)
- 实现IP钉住机制防止验证后的DNS rebinding攻击
- 自定义Transport.DialContext强制只允许拨号到验证过的公网IP
- 扩展IP黑名单,包括云metadata地址(169.254.169.254)
- 添加完整的单元测试覆盖
2. **OpsAlertService生命周期管理** (wire.go)
- 在ProvideOpsMetricsCollector中添加opsAlertService.Start()调用
- 确保stopCtx正确初始化,避免nil指针问题
- 实现防御式启动,保证服务启动顺序
3. **数据库查询排序** (ops_repo.go)
- 在ListRecentSystemMetrics中添加显式ORDER BY updated_at DESC, id DESC
- 在GetLatestSystemMetric中添加排序保证
- 避免数据库返回顺序不确定导致告警误判
### P1 重要问题
4. **并发安全** (ops_metrics_collector.go)
- 为lastGCPauseTotal字段添加sync.Mutex保护
- 防止数据竞争
5. **Goroutine泄漏** (ops_error_logger.go)
- 实现worker pool模式限制并发goroutine数量
- 使用256容量缓冲队列和10个固定worker
- 非阻塞投递,队列满时丢弃任务
6. **生命周期控制** (ops_alert_service.go)
- 添加Start/Stop方法实现优雅关闭
- 使用context控制goroutine生命周期
- 实现WaitGroup等待后台任务完成
7. **Webhook URL验证** (ops_alert_service.go)
- 防止SSRF攻击:验证scheme、禁止内网IP
- DNS解析验证,拒绝解析到私有IP的域名
- 添加8个单元测试覆盖各种攻击场景
8. **资源泄漏** (ops_repo.go)
- 修复多处defer rows.Close()问题
- 简化冗余的defer func()包装
9. **HTTP超时控制** (ops_alert_service.go)
- 创建带10秒超时的http.Client
- 添加buildWebhookHTTPClient辅助函数
- 防止HTTP请求无限期挂起
10. **数据库查询优化** (ops_repo.go)
- 将GetWindowStats的4次独立查询合并为1次CTE查询
- 减少网络往返和表扫描次数
- 显著提升性能
11. **重试机制** (ops_alert_service.go)
- 实现邮件发送重试:最多3次,指数退避(1s/2s/4s)
- 添加webhook备用通道
- 实现完整的错误处理和日志记录
12. **魔法数字** (ops_repo.go, ops_metrics_collector.go)
- 提取硬编码数字为有意义的常量
- 提高代码可读性和可维护性
## 测试验证
- ✅ go test ./internal/service -tags opsalert_unit 通过
- ✅ 所有webhook验证测试通过
- ✅ 重试机制测试通过
## 影响范围
- 运维监控系统安全性显著提升
- 系统稳定性和性能优化
- 无破坏性变更,向后兼容
* feat(ops): 运维监控系统V2 - 完整实现
## 核心功能
- 运维监控仪表盘V2(实时监控、历史趋势、告警管理)
- WebSocket实时QPS/TPS监控(30s心跳,自动重连)
- 系统指标采集(CPU、内存、延迟、错误率等)
- 多维度统计分析(按provider、model、user等维度)
- 告警规则管理(阈值配置、通知渠道)
- 错误日志追踪(详细错误信息、堆栈跟踪)
## 数据库Schema (Migration 025)
### 扩展现有表
- ops_system_metrics: 新增RED指标、错误分类、延迟指标、资源指标、业务指标
- ops_alert_rules: 新增JSONB字段(dimension_filters, notify_channels, notify_config)
### 新增表
- ops_dimension_stats: 多维度统计数据
- ops_data_retention_config: 数据保留策略配置
### 新增视图和函数
- ops_latest_metrics: 最新1分钟窗口指标(已修复字段名和window过滤)
- ops_active_alerts: 当前活跃告警(已修复字段名和状态值)
- calculate_health_score: 健康分数计算函数
## 一致性修复(98/100分)
### P0级别(阻塞Migration)
- ✅ 修复ops_latest_metrics视图字段名(latency_p99→p99_latency_ms, cpu_usage→cpu_usage_percent)
- ✅ 修复ops_active_alerts视图字段名(metric→metric_type, triggered_at→fired_at, trigger_value→metric_value, threshold→threshold_value)
- ✅ 统一告警历史表名(删除ops_alert_history,使用ops_alert_events)
- ✅ 统一API参数限制(ListMetricsHistory和ListErrorLogs的limit改为5000)
### P1级别(功能完整性)
- ✅ 修复ops_latest_metrics视图未过滤window_minutes(添加WHERE m.window_minutes = 1)
- ✅ 修复数据回填UPDATE逻辑(QPS计算改为request_count/(window_minutes*60.0))
- ✅ 添加ops_alert_rules JSONB字段后端支持(Go结构体+序列化)
### P2级别(优化)
- ✅ 前端WebSocket自动重连(指数退避1s→2s→4s→8s→16s,最大5次)
- ✅ 后端WebSocket心跳检测(30s ping,60s pong超时)
## 技术实现
### 后端 (Go)
- Handler层: ops_handler.go(REST API), ops_ws_handler.go(WebSocket)
- Service层: ops_service.go(核心逻辑), ops_cache.go(缓存), ops_alerts.go(告警)
- Repository层: ops_repo.go(数据访问), ops.go(模型定义)
- 路由: admin.go(新增ops相关路由)
- 依赖注入: wire_gen.go(自动生成)
### 前端 (Vue3 + TypeScript)
- 组件: OpsDashboardV2.vue(仪表盘主组件)
- API: ops.ts(REST API + WebSocket封装)
- 路由: index.ts(新增/admin/ops路由)
- 国际化: en.ts, zh.ts(中英文支持)
## 测试验证
- ✅ 所有Go测试通过
- ✅ Migration可正常执行
- ✅ WebSocket连接稳定
- ✅ 前后端数据结构对齐
* refactor: 代码清理和测试优化
## 测试文件优化
- 简化integration test fixtures和断言
- 优化test helper函数
- 统一测试数据格式
## 代码清理
- 移除未使用的代码和注释
- 简化concurrency_cache实现
- 优化middleware错误处理
## 小修复
- 修复gateway_handler和openai_gateway_handler的小问题
- 统一代码风格和格式
变更统计: 27个文件,292行新增,322行删除(净减少30行)
* fix(ops): 运维监控系统安全加固和功能优化
## 安全增强
- feat(security): WebSocket日志脱敏机制,防止token/api_key泄露
- feat(security): X-Forwarded-Host白名单验证,防止CSRF绕过
- feat(security): Origin策略配置化,支持strict/permissive模式
- feat(auth): WebSocket认证支持query参数传递token
## 配置优化
- feat(config): 支持环境变量配置代理信任和Origin策略
- OPS_WS_TRUST_PROXY
- OPS_WS_TRUSTED_PROXIES
- OPS_WS_ORIGIN_POLICY
- fix(ops): 错误日志查询限流从5000降至500,优化内存使用
## 架构改进
- refactor(ops): 告警服务解耦,独立运行评估定时器
- refactor(ops): OpsDashboard统一版本,移除V2分离
## 测试和文档
- test(ops): 添加WebSocket安全验证单元测试(8个测试用例)
- test(ops): 添加告警服务集成测试
- docs(api): 更新API文档,标注限流变更
- docs: 添加CHANGELOG记录breaking changes
## 修复文件
Backend:
- backend/internal/server/middleware/logger.go
- backend/internal/handler/admin/ops_handler.go
- backend/internal/handler/admin/ops_ws_handler.go
- backend/internal/server/middleware/admin_auth.go
- backend/internal/service/ops_alert_service.go
- backend/internal/service/ops_metrics_collector.go
- backend/internal/service/wire.go
Frontend:
- frontend/src/views/admin/ops/OpsDashboard.vue
- frontend/src/router/index.ts
- frontend/src/api/admin/ops.ts
Tests:
- backend/internal/handler/admin/ops_ws_handler_test.go (新增)
- backend/internal/service/ops_alert_service_integration_test.go (新增)
Docs:
- CHANGELOG.md (新增)
- docs/API-运维监控中心2.0.md (更新)
* fix(migrations): 修复calculate_health_score函数类型匹配问题
在ops_latest_metrics视图中添加显式类型转换,确保参数类型与函数签名匹配
* fix(lint): 修复golangci-lint检查发现的所有问题
- 将Redis依赖从service层移到repository层
- 添加错误检查(WebSocket连接和读取超时)
- 运行gofmt格式化代码
- 添加nil指针检查
- 删除未使用的alertService字段
修复问题:
- depguard: 3个(service层不应直接import redis)
- errcheck: 3个(未检查错误返回值)
- gofmt: 2个(代码格式问题)
- staticcheck: 4个(nil指针解引用)
- unused: 1个(未使用字段)
代码统计:
- 修改文件:11个
- 删除代码:490行
- 新增代码:105行
- 净减少:385行
588 lines
21 KiB
Go
588 lines
21 KiB
Go
//go:build integration
|
|
|
|
package repository
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
dbent "github.com/Wei-Shaw/sub2api/ent"
|
|
"github.com/Wei-Shaw/sub2api/ent/accountgroup"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/stretchr/testify/suite"
|
|
)
|
|
|
|
type AccountRepoSuite struct {
|
|
suite.Suite
|
|
ctx context.Context
|
|
client *dbent.Client
|
|
repo *accountRepository
|
|
}
|
|
|
|
func (s *AccountRepoSuite) SetupTest() {
|
|
s.ctx = context.Background()
|
|
tx := testEntTx(s.T())
|
|
s.client = tx.Client()
|
|
s.repo = newAccountRepositoryWithSQL(s.client, tx)
|
|
}
|
|
|
|
func TestAccountRepoSuite(t *testing.T) {
|
|
suite.Run(t, new(AccountRepoSuite))
|
|
}
|
|
|
|
// --- Create / GetByID / Update / Delete ---
|
|
|
|
func (s *AccountRepoSuite) TestCreate() {
|
|
account := &service.Account{
|
|
Name: "test-create",
|
|
Platform: service.PlatformAnthropic,
|
|
Type: service.AccountTypeOAuth,
|
|
Status: service.StatusActive,
|
|
Credentials: map[string]any{},
|
|
Extra: map[string]any{},
|
|
Concurrency: 3,
|
|
Priority: 50,
|
|
Schedulable: true,
|
|
}
|
|
|
|
err := s.repo.Create(s.ctx, account)
|
|
s.Require().NoError(err, "Create")
|
|
s.Require().NotZero(account.ID, "expected ID to be set")
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetByID")
|
|
s.Require().Equal("test-create", got.Name)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestGetByID_NotFound() {
|
|
_, err := s.repo.GetByID(s.ctx, 999999)
|
|
s.Require().Error(err, "expected error for non-existent ID")
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestUpdate() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "original"})
|
|
|
|
account.Name = "updated"
|
|
err := s.repo.Update(s.ctx, account)
|
|
s.Require().NoError(err, "Update")
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetByID after update")
|
|
s.Require().Equal("updated", got.Name)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestDelete() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "to-delete"})
|
|
|
|
err := s.repo.Delete(s.ctx, account.ID)
|
|
s.Require().NoError(err, "Delete")
|
|
|
|
_, err = s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().Error(err, "expected error after delete")
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestDelete_WithGroupBindings() {
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-del"})
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-del"})
|
|
mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
|
|
|
|
err := s.repo.Delete(s.ctx, account.ID)
|
|
s.Require().NoError(err, "Delete should cascade remove bindings")
|
|
|
|
count, err := s.client.AccountGroup.Query().Where(accountgroup.AccountIDEQ(account.ID)).Count(s.ctx)
|
|
s.Require().NoError(err)
|
|
s.Require().Zero(count, "expected bindings to be removed")
|
|
}
|
|
|
|
// --- List / ListWithFilters ---
|
|
|
|
func (s *AccountRepoSuite) TestList() {
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc1"})
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc2"})
|
|
|
|
accounts, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10})
|
|
s.Require().NoError(err, "List")
|
|
s.Require().Len(accounts, 2)
|
|
s.Require().Equal(int64(2), page.Total)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestListWithFilters() {
|
|
tests := []struct {
|
|
name string
|
|
setup func(client *dbent.Client)
|
|
platform string
|
|
accType string
|
|
status string
|
|
search string
|
|
wantCount int
|
|
validate func(accounts []service.Account)
|
|
}{
|
|
{
|
|
name: "filter_by_platform",
|
|
setup: func(client *dbent.Client) {
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic})
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI})
|
|
},
|
|
platform: service.PlatformOpenAI,
|
|
wantCount: 1,
|
|
validate: func(accounts []service.Account) {
|
|
s.Require().Equal(service.PlatformOpenAI, accounts[0].Platform)
|
|
},
|
|
},
|
|
{
|
|
name: "filter_by_type",
|
|
setup: func(client *dbent.Client) {
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "t1", Type: service.AccountTypeOAuth})
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "t2", Type: service.AccountTypeAPIKey})
|
|
},
|
|
accType: service.AccountTypeAPIKey,
|
|
wantCount: 1,
|
|
validate: func(accounts []service.Account) {
|
|
s.Require().Equal(service.AccountTypeAPIKey, accounts[0].Type)
|
|
},
|
|
},
|
|
{
|
|
name: "filter_by_status",
|
|
setup: func(client *dbent.Client) {
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "s1", Status: service.StatusActive})
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "s2", Status: service.StatusDisabled})
|
|
},
|
|
status: service.StatusDisabled,
|
|
wantCount: 1,
|
|
validate: func(accounts []service.Account) {
|
|
s.Require().Equal(service.StatusDisabled, accounts[0].Status)
|
|
},
|
|
},
|
|
{
|
|
name: "filter_by_search",
|
|
setup: func(client *dbent.Client) {
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "alpha-account"})
|
|
mustCreateAccount(s.T(), client, &service.Account{Name: "beta-account"})
|
|
},
|
|
search: "alpha",
|
|
wantCount: 1,
|
|
validate: func(accounts []service.Account) {
|
|
s.Require().Contains(accounts[0].Name, "alpha")
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
s.Run(tt.name, func() {
|
|
// 每个 case 重新获取隔离资源
|
|
tx := testEntTx(s.T())
|
|
client := tx.Client()
|
|
repo := newAccountRepositoryWithSQL(client, tx)
|
|
ctx := context.Background()
|
|
|
|
tt.setup(client)
|
|
|
|
accounts, _, err := repo.ListWithFilters(ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, tt.platform, tt.accType, tt.status, tt.search)
|
|
s.Require().NoError(err)
|
|
s.Require().Len(accounts, tt.wantCount)
|
|
if tt.validate != nil {
|
|
tt.validate(accounts)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// --- ListByGroup / ListActive / ListByPlatform ---
|
|
|
|
func (s *AccountRepoSuite) TestListByGroup() {
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-list"})
|
|
acc1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Status: service.StatusActive})
|
|
acc2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Status: service.StatusActive})
|
|
mustBindAccountToGroup(s.T(), s.client, acc1.ID, group.ID, 2)
|
|
mustBindAccountToGroup(s.T(), s.client, acc2.ID, group.ID, 1)
|
|
|
|
accounts, err := s.repo.ListByGroup(s.ctx, group.ID)
|
|
s.Require().NoError(err, "ListByGroup")
|
|
s.Require().Len(accounts, 2)
|
|
// Should be ordered by priority
|
|
s.Require().Equal(acc2.ID, accounts[0].ID, "expected acc2 first (priority=1)")
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestListActive() {
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "active1", Status: service.StatusActive})
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "inactive1", Status: service.StatusDisabled})
|
|
|
|
accounts, err := s.repo.ListActive(s.ctx)
|
|
s.Require().NoError(err, "ListActive")
|
|
s.Require().Len(accounts, 1)
|
|
s.Require().Equal("active1", accounts[0].Name)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestListByPlatform() {
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "p1", Platform: service.PlatformAnthropic, Status: service.StatusActive})
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "p2", Platform: service.PlatformOpenAI, Status: service.StatusActive})
|
|
|
|
accounts, err := s.repo.ListByPlatform(s.ctx, service.PlatformAnthropic)
|
|
s.Require().NoError(err, "ListByPlatform")
|
|
s.Require().Len(accounts, 1)
|
|
s.Require().Equal(service.PlatformAnthropic, accounts[0].Platform)
|
|
}
|
|
|
|
// --- Preload and VirtualFields ---
|
|
|
|
func (s *AccountRepoSuite) TestPreload_And_VirtualFields() {
|
|
proxy := mustCreateProxy(s.T(), s.client, &service.Proxy{Name: "p1"})
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g1"})
|
|
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{
|
|
Name: "acc1",
|
|
ProxyID: &proxy.ID,
|
|
})
|
|
mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetByID")
|
|
s.Require().NotNil(got.Proxy, "expected Proxy preload")
|
|
s.Require().Equal(proxy.ID, got.Proxy.ID)
|
|
s.Require().Len(got.GroupIDs, 1, "expected GroupIDs to be populated")
|
|
s.Require().Equal(group.ID, got.GroupIDs[0])
|
|
s.Require().Len(got.Groups, 1, "expected Groups to be populated")
|
|
s.Require().Equal(group.ID, got.Groups[0].ID)
|
|
|
|
accounts, page, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, "", "", "", "acc")
|
|
s.Require().NoError(err, "ListWithFilters")
|
|
s.Require().Equal(int64(1), page.Total)
|
|
s.Require().Len(accounts, 1)
|
|
s.Require().NotNil(accounts[0].Proxy, "expected Proxy preload in list")
|
|
s.Require().Equal(proxy.ID, accounts[0].Proxy.ID)
|
|
s.Require().Len(accounts[0].GroupIDs, 1, "expected GroupIDs in list")
|
|
s.Require().Equal(group.ID, accounts[0].GroupIDs[0])
|
|
}
|
|
|
|
// --- GroupBinding / AddToGroup / RemoveFromGroup / BindGroups / GetGroups ---
|
|
|
|
func (s *AccountRepoSuite) TestGroupBinding_And_BindGroups() {
|
|
g1 := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g1"})
|
|
g2 := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g2"})
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc"})
|
|
|
|
s.Require().NoError(s.repo.AddToGroup(s.ctx, account.ID, g1.ID, 10), "AddToGroup")
|
|
groups, err := s.repo.GetGroups(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetGroups")
|
|
s.Require().Len(groups, 1, "expected 1 group")
|
|
s.Require().Equal(g1.ID, groups[0].ID)
|
|
|
|
s.Require().NoError(s.repo.RemoveFromGroup(s.ctx, account.ID, g1.ID), "RemoveFromGroup")
|
|
groups, err = s.repo.GetGroups(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetGroups after remove")
|
|
s.Require().Empty(groups, "expected 0 groups after remove")
|
|
|
|
s.Require().NoError(s.repo.BindGroups(s.ctx, account.ID, []int64{g1.ID, g2.ID}), "BindGroups")
|
|
groups, err = s.repo.GetGroups(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetGroups after bind")
|
|
s.Require().Len(groups, 2, "expected 2 groups after bind")
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestBindGroups_EmptyList() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-empty"})
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-empty"})
|
|
mustBindAccountToGroup(s.T(), s.client, account.ID, group.ID, 1)
|
|
|
|
s.Require().NoError(s.repo.BindGroups(s.ctx, account.ID, []int64{}), "BindGroups empty")
|
|
|
|
groups, err := s.repo.GetGroups(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().Empty(groups, "expected 0 groups after binding empty list")
|
|
}
|
|
|
|
// --- Schedulable ---
|
|
|
|
func (s *AccountRepoSuite) TestListSchedulable() {
|
|
now := time.Now()
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sched"})
|
|
|
|
okAcc := mustCreateAccount(s.T(), s.client, &service.Account{Name: "ok", Schedulable: true})
|
|
mustBindAccountToGroup(s.T(), s.client, okAcc.ID, group.ID, 1)
|
|
|
|
future := now.Add(10 * time.Minute)
|
|
overloaded := mustCreateAccount(s.T(), s.client, &service.Account{Name: "over", Schedulable: true, OverloadUntil: &future})
|
|
mustBindAccountToGroup(s.T(), s.client, overloaded.ID, group.ID, 1)
|
|
|
|
sched, err := s.repo.ListSchedulable(s.ctx)
|
|
s.Require().NoError(err, "ListSchedulable")
|
|
ids := idsOfAccounts(sched)
|
|
s.Require().Contains(ids, okAcc.ID)
|
|
s.Require().NotContains(ids, overloaded.ID)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestListSchedulableByGroupID_TimeBoundaries_And_StatusUpdates() {
|
|
now := time.Now()
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sched"})
|
|
|
|
okAcc := mustCreateAccount(s.T(), s.client, &service.Account{Name: "ok", Schedulable: true})
|
|
mustBindAccountToGroup(s.T(), s.client, okAcc.ID, group.ID, 1)
|
|
|
|
future := now.Add(10 * time.Minute)
|
|
overloaded := mustCreateAccount(s.T(), s.client, &service.Account{Name: "over", Schedulable: true, OverloadUntil: &future})
|
|
mustBindAccountToGroup(s.T(), s.client, overloaded.ID, group.ID, 1)
|
|
|
|
rateLimited := mustCreateAccount(s.T(), s.client, &service.Account{Name: "rl", Schedulable: true})
|
|
mustBindAccountToGroup(s.T(), s.client, rateLimited.ID, group.ID, 1)
|
|
s.Require().NoError(s.repo.SetRateLimited(s.ctx, rateLimited.ID, now.Add(10*time.Minute)), "SetRateLimited")
|
|
|
|
s.Require().NoError(s.repo.SetError(s.ctx, overloaded.ID, "boom"), "SetError")
|
|
|
|
sched, err := s.repo.ListSchedulableByGroupID(s.ctx, group.ID)
|
|
s.Require().NoError(err, "ListSchedulableByGroupID")
|
|
s.Require().Len(sched, 1, "expected only ok account schedulable")
|
|
s.Require().Equal(okAcc.ID, sched[0].ID)
|
|
|
|
s.Require().NoError(s.repo.ClearRateLimit(s.ctx, rateLimited.ID), "ClearRateLimit")
|
|
sched2, err := s.repo.ListSchedulableByGroupID(s.ctx, group.ID)
|
|
s.Require().NoError(err, "ListSchedulableByGroupID after ClearRateLimit")
|
|
s.Require().Len(sched2, 2, "expected 2 schedulable accounts after ClearRateLimit")
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestListSchedulableByPlatform() {
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic, Schedulable: true})
|
|
mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI, Schedulable: true})
|
|
|
|
accounts, err := s.repo.ListSchedulableByPlatform(s.ctx, service.PlatformAnthropic)
|
|
s.Require().NoError(err)
|
|
s.Require().Len(accounts, 1)
|
|
s.Require().Equal(service.PlatformAnthropic, accounts[0].Platform)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestListSchedulableByGroupIDAndPlatform() {
|
|
group := mustCreateGroup(s.T(), s.client, &service.Group{Name: "g-sp"})
|
|
a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a1", Platform: service.PlatformAnthropic, Schedulable: true})
|
|
a2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "a2", Platform: service.PlatformOpenAI, Schedulable: true})
|
|
mustBindAccountToGroup(s.T(), s.client, a1.ID, group.ID, 1)
|
|
mustBindAccountToGroup(s.T(), s.client, a2.ID, group.ID, 2)
|
|
|
|
accounts, err := s.repo.ListSchedulableByGroupIDAndPlatform(s.ctx, group.ID, service.PlatformAnthropic)
|
|
s.Require().NoError(err)
|
|
s.Require().Len(accounts, 1)
|
|
s.Require().Equal(a1.ID, accounts[0].ID)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestSetSchedulable() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-sched", Schedulable: true})
|
|
|
|
s.Require().NoError(s.repo.SetSchedulable(s.ctx, account.ID, false))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().False(got.Schedulable)
|
|
}
|
|
|
|
// --- SetOverloaded / SetRateLimited / ClearRateLimit ---
|
|
|
|
func (s *AccountRepoSuite) TestSetOverloaded() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-over"})
|
|
until := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
|
|
|
|
s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().NotNil(got.OverloadUntil)
|
|
s.Require().WithinDuration(until, *got.OverloadUntil, time.Second)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestSetRateLimited() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl"})
|
|
resetAt := time.Date(2025, 6, 15, 14, 0, 0, 0, time.UTC)
|
|
|
|
s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, resetAt))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().NotNil(got.RateLimitedAt)
|
|
s.Require().NotNil(got.RateLimitResetAt)
|
|
s.Require().WithinDuration(resetAt, *got.RateLimitResetAt, time.Second)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestClearRateLimit() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"})
|
|
until := time.Now().Add(1 * time.Hour)
|
|
s.Require().NoError(s.repo.SetOverloaded(s.ctx, account.ID, until))
|
|
s.Require().NoError(s.repo.SetRateLimited(s.ctx, account.ID, until))
|
|
|
|
s.Require().NoError(s.repo.ClearRateLimit(s.ctx, account.ID))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().Nil(got.RateLimitedAt)
|
|
s.Require().Nil(got.RateLimitResetAt)
|
|
s.Require().Nil(got.OverloadUntil)
|
|
}
|
|
|
|
// --- UpdateLastUsed ---
|
|
|
|
func (s *AccountRepoSuite) TestUpdateLastUsed() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-used"})
|
|
s.Require().Nil(account.LastUsedAt)
|
|
|
|
s.Require().NoError(s.repo.UpdateLastUsed(s.ctx, account.ID))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().NotNil(got.LastUsedAt)
|
|
}
|
|
|
|
// --- SetError ---
|
|
|
|
func (s *AccountRepoSuite) TestSetError() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-err", Status: service.StatusActive})
|
|
|
|
s.Require().NoError(s.repo.SetError(s.ctx, account.ID, "something went wrong"))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().Equal(service.StatusError, got.Status)
|
|
s.Require().Equal("something went wrong", got.ErrorMessage)
|
|
}
|
|
|
|
// --- UpdateSessionWindow ---
|
|
|
|
func (s *AccountRepoSuite) TestUpdateSessionWindow() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-win"})
|
|
start := time.Date(2025, 6, 15, 10, 0, 0, 0, time.UTC)
|
|
end := time.Date(2025, 6, 15, 15, 0, 0, 0, time.UTC)
|
|
|
|
s.Require().NoError(s.repo.UpdateSessionWindow(s.ctx, account.ID, &start, &end, "active"))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().NotNil(got.SessionWindowStart)
|
|
s.Require().NotNil(got.SessionWindowEnd)
|
|
s.Require().Equal("active", got.SessionWindowStatus)
|
|
}
|
|
|
|
// --- UpdateExtra ---
|
|
|
|
func (s *AccountRepoSuite) TestUpdateExtra_MergesFields() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{
|
|
Name: "acc-extra",
|
|
Extra: map[string]any{"a": "1"},
|
|
})
|
|
s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{"b": "2"}), "UpdateExtra")
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err, "GetByID")
|
|
s.Require().Equal("1", got.Extra["a"])
|
|
s.Require().Equal("2", got.Extra["b"])
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestUpdateExtra_EmptyUpdates() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-extra-empty"})
|
|
s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{}))
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestUpdateExtra_NilExtra() {
|
|
account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-nil-extra", Extra: nil})
|
|
s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{"key": "val"}))
|
|
|
|
got, err := s.repo.GetByID(s.ctx, account.ID)
|
|
s.Require().NoError(err)
|
|
s.Require().Equal("val", got.Extra["key"])
|
|
}
|
|
|
|
// --- GetByCRSAccountID ---
|
|
|
|
func (s *AccountRepoSuite) TestGetByCRSAccountID() {
|
|
crsID := "crs-12345"
|
|
mustCreateAccount(s.T(), s.client, &service.Account{
|
|
Name: "acc-crs",
|
|
Extra: map[string]any{"crs_account_id": crsID},
|
|
})
|
|
|
|
got, err := s.repo.GetByCRSAccountID(s.ctx, crsID)
|
|
s.Require().NoError(err)
|
|
s.Require().NotNil(got)
|
|
s.Require().Equal("acc-crs", got.Name)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestGetByCRSAccountID_NotFound() {
|
|
got, err := s.repo.GetByCRSAccountID(s.ctx, "non-existent")
|
|
s.Require().NoError(err)
|
|
s.Require().Nil(got)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestGetByCRSAccountID_EmptyString() {
|
|
got, err := s.repo.GetByCRSAccountID(s.ctx, "")
|
|
s.Require().NoError(err)
|
|
s.Require().Nil(got)
|
|
}
|
|
|
|
// --- BulkUpdate ---
|
|
|
|
func (s *AccountRepoSuite) TestBulkUpdate() {
|
|
a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk1", Priority: 1})
|
|
a2 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk2", Priority: 1})
|
|
|
|
newPriority := 99
|
|
affected, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID, a2.ID}, service.AccountBulkUpdate{
|
|
Priority: &newPriority,
|
|
})
|
|
s.Require().NoError(err)
|
|
s.Require().GreaterOrEqual(affected, int64(1), "expected at least one affected row")
|
|
|
|
got1, _ := s.repo.GetByID(s.ctx, a1.ID)
|
|
got2, _ := s.repo.GetByID(s.ctx, a2.ID)
|
|
s.Require().Equal(99, got1.Priority)
|
|
s.Require().Equal(99, got2.Priority)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestBulkUpdate_MergeCredentials() {
|
|
a1 := mustCreateAccount(s.T(), s.client, &service.Account{
|
|
Name: "bulk-cred",
|
|
Credentials: map[string]any{"existing": "value"},
|
|
})
|
|
|
|
_, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{
|
|
Credentials: map[string]any{"new_key": "new_value"},
|
|
})
|
|
s.Require().NoError(err)
|
|
|
|
got, _ := s.repo.GetByID(s.ctx, a1.ID)
|
|
s.Require().Equal("value", got.Credentials["existing"])
|
|
s.Require().Equal("new_value", got.Credentials["new_key"])
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestBulkUpdate_MergeExtra() {
|
|
a1 := mustCreateAccount(s.T(), s.client, &service.Account{
|
|
Name: "bulk-extra",
|
|
Extra: map[string]any{"existing": "val"},
|
|
})
|
|
|
|
_, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{
|
|
Extra: map[string]any{"new_key": "new_val"},
|
|
})
|
|
s.Require().NoError(err)
|
|
|
|
got, _ := s.repo.GetByID(s.ctx, a1.ID)
|
|
s.Require().Equal("val", got.Extra["existing"])
|
|
s.Require().Equal("new_val", got.Extra["new_key"])
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestBulkUpdate_EmptyIDs() {
|
|
affected, err := s.repo.BulkUpdate(s.ctx, []int64{}, service.AccountBulkUpdate{})
|
|
s.Require().NoError(err)
|
|
s.Require().Zero(affected)
|
|
}
|
|
|
|
func (s *AccountRepoSuite) TestBulkUpdate_EmptyUpdates() {
|
|
a1 := mustCreateAccount(s.T(), s.client, &service.Account{Name: "bulk-empty"})
|
|
|
|
affected, err := s.repo.BulkUpdate(s.ctx, []int64{a1.ID}, service.AccountBulkUpdate{})
|
|
s.Require().NoError(err)
|
|
s.Require().Zero(affected)
|
|
}
|
|
|
|
func idsOfAccounts(accounts []service.Account) []int64 {
|
|
out := make([]int64, 0, len(accounts))
|
|
for i := range accounts {
|
|
out = append(out, accounts[i].ID)
|
|
}
|
|
return out
|
|
}
|