LokiStack/backend/services/daily_volume_test.go

package services

import (
	"testing"
	"time"

	"lokistack/models"
	"lokistack/repository"
)

func setupDailyVolumeDB(t *testing.T) {
	t.Helper()
	t.Setenv("DATA_DIR", t.TempDir())
	t.Setenv("JWT_SECRET", "test-jwt-secret-with-at-least-32-characters")
	t.Setenv("ENCRYPTION_KEY", "test-encryption-key-with-at-least-32-characters")
	if err := repository.InitDB(); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = repository.Shutdown() })

	// 每天的快照都挂着数据源，prune 会清掉"数据源已删除"的孤儿快照，
	// 所以测试里先登记好 id=1 的数据源。
	db := repository.GetDB()
	db.Lock()
	if len(db.LogSources) == 0 {
		db.LogSources = append(db.LogSources, models.LogSource{ID: 1, Name: "test-loki", TargetType: "url"})
	}
	db.Unlock()
}

func snapshotsByDate(t *testing.T) map[string]models.DailyVolumeSnapshot {
	t.Helper()
	got := map[string]models.DailyVolumeSnapshot{}
	db := repository.GetDB()
	db.RLock()
	defer db.RUnlock()
	for _, snap := range db.DailyVolumeSnapshots {
		got[snap.Date] = snap
	}
	return got
}

// 过去的天不可能"还在进行中"：只有在那天结束之后才抓的快照才算完整一天，
// 当天中途抓到的必须保持 partial —— 否则会把"只抓到半天"伪装成完整数据。
func TestPruneDailyVolumeNormalizesPartialAndDropsExpired(t *testing.T) {
	setupDailyVolumeDB(t)
	now := time.Now()
	midday := func(d time.Time) time.Time {
		return time.Date(d.Year(), d.Month(), d.Day(), 12, 0, 0, 0, time.Local)
	}
	yesterday := now.AddDate(0, 0, -1)
	dayBefore := now.AddDate(0, 0, -2)
	expired := now.AddDate(0, 0, -(dailyVolumeRetentionDays + 5))

	db := repository.GetDB()
	db.Lock()
	db.DailyVolumeSnapshots = []models.DailyVolumeSnapshot{
		// 昨天结束之后（次日 12:00）抓的 → 完整一天
		{ID: 1, Date: dateKey(yesterday), SourceID: 1, Bytes: 100, Entries: 7,
			Partial: true, CapturedAt: midday(now)},
		// 前天中午抓的 → 只有半天，保持 partial
		{ID: 2, Date: dateKey(dayBefore), SourceID: 1, Bytes: 50,
			Partial: true, CapturedAt: midday(dayBefore)},
		// 今天 → 永远 partial
		{ID: 3, Date: dateKey(now), SourceID: 1, Bytes: 10, Partial: true},
		// 超过留存期 → 必须被清理
		{ID: 4, Date: dateKey(expired), SourceID: 1, Bytes: 999},
		// 数据源已删除的孤儿快照 → 必须被清理
		{ID: 5, Date: dateKey(yesterday), SourceID: 99, Bytes: 888},
	}
	db.Unlock()

	pruneDailyVolumeSnapshots()

	got := snapshotsByDate(t)
	if _, ok := got[dateKey(expired)]; ok {
		t.Error("超过留存期的快照没有被清理")
	}
	// 孤儿快照（SourceID=99 的数据源已不存在）必须被清掉
	db = repository.GetDB()
	db.RLock()
	for _, snap := range db.DailyVolumeSnapshots {
		if snap.SourceID == 99 {
			t.Error("数据源已删除的孤儿快照没有被清理")
		}
	}
	db.RUnlock()
	if snap := got[dateKey(yesterday)]; snap.Partial {
		t.Error("昨天结束之后抓的快照应标记为完整（partial=false）")
	}
	if snap := got[dateKey(dayBefore)]; !snap.Partial {
		t.Error("当天中途抓的快照必须保持 partial=true")
	}
	if snap := got[dateKey(now)]; !snap.Partial {
		t.Error("今天必须保持 partial=true")
	}
	if snap := got[dateKey(yesterday)]; snap.Entries != 7 {
		t.Errorf("条数不应被清理流程清空，得到 %d", snap.Entries)
	}
}

// 刷新体积不能把已经抓到的条数（贵、每天最多 3 次）冲掉。
func TestUpsertVolumeSnapshotKeepsEntries(t *testing.T) {
	setupDailyVolumeDB(t)

	db := repository.GetDB()
	db.Lock()
	db.DailyVolumeSnapshots = []models.DailyVolumeSnapshot{
		{ID: 1, Date: "2026-10-01", SourceID: 1, Bytes: 111, Entries: 42, Partial: true},
	}
	db.Unlock()

	if err := upsertVolumeSnapshot(1, "2026-10-01", 1000, false); err != nil {
		t.Fatal(err)
	}
	got := snapshotsByDate(t)["2026-10-01"]
	if got.Bytes != 1000 || got.Entries != 42 || got.Partial {
		t.Fatalf("更新后 bytes=%d entries=%d partial=%v；期望 1000/42/false", got.Bytes, got.Entries, got.Partial)
	}

	if err := upsertVolumeSnapshot(1, "2026-10-02", 2000, true); err != nil {
		t.Fatal(err)
	}
	if got := snapshotsByDate(t); len(got) != 2 {
		t.Fatalf("新增一天后应有 2 条快照，得到 %d", len(got))
	}
}

// 接口必须始终返回 days 个点，缺的标 Missing（"不知道"），绝不伪造 0（"确实没日志"）。
func TestGetDailyVolumeHistoryReturnsAllDaysWithMissing(t *testing.T) {
	setupDailyVolumeDB(t)
	now := time.Now()

	db := repository.GetDB()
	db.Lock()
	db.DailyVolumeSnapshots = []models.DailyVolumeSnapshot{
		{ID: 1, Date: dateKey(now), SourceID: 1, Bytes: 123, Partial: true},
		{ID: 2, Date: dateKey(now), SourceID: 2, Bytes: 999}, // 别的数据源，不应混进来
	}
	db.Unlock()

	points := GetDailyVolumeHistory(1, 7)
	if len(points) != 7 {
		t.Fatalf("期望 7 个点，得到 %d", len(points))
	}
	today := points[len(points)-1]
	if today.Missing || today.Bytes != 123 || !today.Partial {
		t.Errorf("今天应命中快照：missing=%v bytes=%d partial=%v", today.Missing, today.Bytes, today.Partial)
	}
	for i := 0; i < len(points)-1; i++ {
		if !points[i].Missing || points[i].Bytes != 0 {
			t.Errorf("%s 应标记 Missing 且 bytes=0，得到 missing=%v bytes=%d",
				points[i].Date, points[i].Missing, points[i].Bytes)
		}
	}
}