sub2api/backend/internal/repository/channel_monitor_quota_integration_test.go

//go:build integration

package repository

import (
	"context"
	"testing"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/domain"
	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/stretchr/testify/require"
)

// 配额模式 repo 层集成测试：
//   - Create/GetByID/Update 的 check_mode/account_id 往返
//   - InsertHistoryBatch → ListHistory / ListLatestForMonitorIDs 的 quota JSONB 回读
//     （裸 SQL 列 + scanMonitorQuota），以及探活模式旧行 quota=NULL 的兼容
//
// 注意 channelMonitorRepository 的 GetByID/裸 SQL 走全局 client（不识别 tx ctx），
// 因此本文件用 integrationEntClient 直连 + t.Cleanup 显式清理，不走 testEntTx 回滚。

func TestChannelMonitorQuotaModeRoundTrip(t *testing.T) {
	ctx := context.Background()
	repo := NewChannelMonitorRepository(integrationEntClient, integrationDB)

	account := mustCreateAccount(t, integrationEntClient, &service.Account{
		Name: "quota-linked-kimi", Platform: domain.PlatformKimi, Type: service.AccountTypeAPIKey,
		Credentials: map[string]any{"api_key": "sk-kimi", "account_mode": service.AccountModeCoding},
	})
	t.Cleanup(func() {
		_ = integrationEntClient.Account.DeleteOneID(account.ID).Exec(ctx)
	})

	created := &service.ChannelMonitor{
		Name:             "kimi-quota-roundtrip",
		Provider:         service.MonitorProviderKimi,
		APIMode:          service.MonitorAPIModeChatCompletions,
		Endpoint:         "",
		APIKey:           "encrypted-empty",
		PrimaryModel:     "quota",
		Enabled:          true,
		IntervalSeconds:  60,
		CheckMode:        service.MonitorCheckModeQuota,
		AccountID:        &account.ID,
		BodyOverrideMode: service.MonitorBodyOverrideModeOff,
	}
	require.NoError(t, repo.Create(ctx, created))
	t.Cleanup(func() {
		_ = repo.Delete(ctx, created.ID)
	})

	loaded, err := repo.GetByID(ctx, created.ID)
	require.NoError(t, err)
	require.Equal(t, service.MonitorCheckModeQuota, loaded.CheckMode)
	require.NotNil(t, loaded.AccountID)
	require.Equal(t, account.ID, *loaded.AccountID)

	// Update：切换模式并清空关联账号（probe 化）。
	loaded.CheckMode = service.MonitorCheckModeProbe
	loaded.AccountID = nil
	loaded.Endpoint = "https://api.moonshot.cn"
	require.NoError(t, repo.Update(ctx, loaded))

	reloaded, err := repo.GetByID(ctx, created.ID)
	require.NoError(t, err)
	require.Equal(t, service.MonitorCheckModeProbe, reloaded.CheckMode)
	require.Nil(t, reloaded.AccountID)

	// 重新绑定账号。
	reloaded.CheckMode = service.MonitorCheckModeQuotaProbe
	reloaded.AccountID = &account.ID
	require.NoError(t, repo.Update(ctx, reloaded))
	final, err := repo.GetByID(ctx, created.ID)
	require.NoError(t, err)
	require.Equal(t, service.MonitorCheckModeQuotaProbe, final.CheckMode)
	require.NotNil(t, final.AccountID)
}

func TestChannelMonitorHistoryQuotaRoundTrip(t *testing.T) {
	ctx := context.Background()
	repo := NewChannelMonitorRepository(integrationEntClient, integrationDB)

	monitor := &service.ChannelMonitor{
		Name:             "quota-history-roundtrip",
		Provider:         service.MonitorProviderOpenAI,
		APIMode:          service.MonitorAPIModeChatCompletions,
		Endpoint:         "https://api.openai.com",
		APIKey:           "encrypted",
		PrimaryModel:     "gpt-test",
		ExtraModels:      []string{"gpt-extra"},
		Enabled:          true,
		IntervalSeconds:  60,
		BodyOverrideMode: service.MonitorBodyOverrideModeOff,
	}
	require.NoError(t, repo.Create(ctx, monitor))
	t.Cleanup(func() {
		_ = repo.Delete(ctx, monitor.ID) // histories 级联删除
	})

	now := time.Now().UTC()
	rows := []*service.ChannelMonitorHistoryRow{
		{
			MonitorID: monitor.ID, Model: "gpt-test", Status: service.MonitorStatusOperational,
			Message: "ok", CheckedAt: now,
			Quota: &domain.MonitorQuotaSnapshot{
				Source:    "usage",
				Success:   true,
				PlanLevel: "PRO",
				Tiers: []domain.MonitorQuotaTier{
					{Window: "5h", UsedPercent: 42.5, Used: 17, Limit: 40, ResetAt: now.Add(time.Hour).Format(time.RFC3339)},
				},
				FetchedAt: now,
			},
		},
		{
			// 探活模式旧行：无 quota（NULL 兼容）。
			MonitorID: monitor.ID, Model: "gpt-extra", Status: service.MonitorStatusOperational,
			Message: "ok", CheckedAt: now,
		},
	}
	require.NoError(t, repo.InsertHistoryBatch(ctx, rows))

	history, err := repo.ListHistory(ctx, monitor.ID, "", 10)
	require.NoError(t, err)
	require.Len(t, history, 2)

	byModel := map[string]*service.ChannelMonitorHistoryEntry{}
	for _, entry := range history {
		byModel[entry.Model] = entry
	}
	withQuota := byModel["gpt-test"]
	require.NotNil(t, withQuota.Quota)
	require.True(t, withQuota.Quota.Success)
	require.Equal(t, "usage", withQuota.Quota.Source)
	require.Equal(t, "PRO", withQuota.Quota.PlanLevel)
	require.Len(t, withQuota.Quota.Tiers, 1)
	require.Equal(t, "5h", withQuota.Quota.Tiers[0].Window)
	require.InDelta(t, 42.5, withQuota.Quota.Tiers[0].UsedPercent, 0.001)
	require.Nil(t, byModel["gpt-extra"].Quota, "probe rows must read back as NULL quota")

	// 用户视图聚合：主模型最近一行带快照。
	latest, err := repo.ListLatestForMonitorIDs(ctx, []int64{monitor.ID})
	require.NoError(t, err)
	primaryRows := latest[monitor.ID]
	require.NotEmpty(t, primaryRows)
	var primaryQuota *domain.MonitorQuotaSnapshot
	for _, row := range primaryRows {
		if row.Model == "gpt-test" {
			primaryQuota = row.Quota
		}
	}
	require.NotNil(t, primaryQuota, "ListLatestForMonitorIDs must surface quota for the primary model")
	require.True(t, primaryQuota.Success)
}