// Copyright 2026 tamakyi. All rights reserved. // Use of this source code is governed by a MIT-style // license that can be found in the LICENSE file. package db import ( "context" "testing" "time" "github.com/stretchr/testify/require" ) // TestIPGeoStore 覆盖地理定位缓存的读写、批量取、覆盖与过期清理。 // // 这张表是配额的保护层:同一批攻击者会在多份报告里反复出现, // 「查一次、之后全靠缓存」能不能成立全看这里的语义对不对。 func TestIPGeoStore(t *testing.T) { database, cleanup := newTestDB(t) ctx := context.Background() t.Cleanup(func() { if err := cleanup("ip_geo_cache"); err != nil { t.Fatal(err) } }) store := NewIPGeoStore(database) now := time.Now() records := []*IPGeoRecord{ { IPKey: "203.0.113.1", Country: "美国", CountryCode: "US", Region: "Virginia", City: "Ashburn", ISP: "Amazon", ASN: "AS14618", Lat: 39.04, Lon: -77.49, HasCoords: true, Source: "ipapi", ResolvedAt: now, }, { IPKey: "2001:db8::/64", Country: "中国", CountryCode: "CN", Region: "北京市", City: "北京", Lat: 39.9, Lon: 116.4, HasCoords: true, Source: "tencent", ResolvedAt: now, }, { // 查不到:Failed 标记要能落库,否则每次报告都会重试同一批地址 IPKey: "198.51.100.9", Failed: true, ResolvedAt: now, }, } require.NoError(t, store.UpsertMany(ctx, records)) got, err := store.GetMany(ctx, []string{"203.0.113.1", "2001:db8::/64", "198.51.100.9", "10.0.0.1"}) require.NoError(t, err) require.Len(t, got, 3, "不存在的键不应出现在结果里") require.Equal(t, "Ashburn", got["203.0.113.1"].City) require.Equal(t, "AS14618", got["203.0.113.1"].ASN) require.True(t, got["203.0.113.1"].HasCoords) require.Equal(t, "tencent", got["2001:db8::/64"].Source) require.True(t, got["198.51.100.9"].Failed) // 覆盖写:同一个键再写一次应当更新而不是报主键冲突 require.NoError(t, store.UpsertMany(ctx, []*IPGeoRecord{{ IPKey: "203.0.113.1", Country: "美国", CountryCode: "US", City: "Reston", Lat: 38.96, Lon: -77.35, HasCoords: true, Source: "tencent", ResolvedAt: now.Add(time.Hour), }})) got2, err := store.GetMany(ctx, []string{"203.0.113.1"}) require.NoError(t, err) require.Equal(t, "Reston", got2["203.0.113.1"].City) require.Equal(t, "tencent", got2["203.0.113.1"].Source) n, err := store.Count(ctx) require.NoError(t, err) require.Equal(t, int64(3), n) } // TestIPGeoStorePurge 覆盖两套 TTL 分开清理:成功留 90 天、失败留 3 天。 func TestIPGeoStorePurge(t *testing.T) { database, cleanup := newTestDB(t) ctx := context.Background() t.Cleanup(func() { if err := cleanup("ip_geo_cache"); err != nil { t.Fatal(err) } }) store := NewIPGeoStore(database) now := time.Now() require.NoError(t, store.UpsertMany(ctx, []*IPGeoRecord{ {IPKey: "1.0.0.1", Country: "澳大利亚", ResolvedAt: now.AddDate(0, 0, -100)}, // 成功且过期 {IPKey: "1.0.0.2", Country: "日本", ResolvedAt: now.AddDate(0, 0, -10)}, // 成功且新鲜 {IPKey: "1.0.0.3", Failed: true, ResolvedAt: now.AddDate(0, 0, -4)}, // 失败且过期 {IPKey: "1.0.0.4", Failed: true, ResolvedAt: now}, // 失败且新鲜 })) // 清失败的旧记录:只删 1.0.0.3 n, err := store.PurgeResolvedBefore(ctx, now.AddDate(0, 0, -3), true) require.NoError(t, err) require.Equal(t, int64(1), n) // 清成功的旧记录:只删 1.0.0.1(不能连失败的一起删,两者 TTL 不同) n, err = store.PurgeResolvedBefore(ctx, now.AddDate(0, 0, -90), false) require.NoError(t, err) require.Equal(t, int64(1), n) left, err := store.GetMany(ctx, []string{"1.0.0.1", "1.0.0.2", "1.0.0.3", "1.0.0.4"}) require.NoError(t, err) require.Len(t, left, 2) require.Contains(t, left, "1.0.0.2") require.Contains(t, left, "1.0.0.4") } // TestSecurityReportStore 覆盖报告台账的幂等语义。 // // 这是「重启不会重复发同一期周报」的唯一保证:period_key 唯一索引 + upsert。 func TestSecurityReportStore(t *testing.T) { database, cleanup := newTestDB(t) ctx := context.Background() t.Cleanup(func() { if err := cleanup("security_report_logs"); err != nil { t.Fatal(err) } }) store := NewSecurityReportStore(database) start := time.Date(2026, 10, 19, 0, 0, 0, 0, time.Local) end := start.AddDate(0, 0, 7) // 首次:不存在 rec, err := store.Get(ctx, "week:2026-W43") require.NoError(t, err) require.Nil(t, rec) require.NoError(t, store.Save(ctx, &SecurityReportLog{ PeriodKey: "week:2026-W43", Kind: "week", WindowStart: start, WindowEnd: end, Hits: 42, UniqueSources: 7, Countries: 3, Recipient: "ops@example.com", Attempts: 1, Error: "smtp timeout", })) rec, err = store.Get(ctx, "week:2026-W43") require.NoError(t, err) require.NotNil(t, rec) require.False(t, rec.Sent) require.Equal(t, 1, rec.Attempts) // 第二次:同一 period_key 再写应当覆盖,而不是插入第二行 —— // 否则「这一期发过没」会被多行记录搞乱。 require.NoError(t, store.Save(ctx, &SecurityReportLog{ PeriodKey: "week:2026-W43", Kind: "week", WindowStart: start, WindowEnd: end, Hits: 42, UniqueSources: 7, Countries: 3, Recipient: "ops@example.com", Attempts: 2, Sent: true, })) rec, err = store.Get(ctx, "week:2026-W43") require.NoError(t, err) require.True(t, rec.Sent) require.Equal(t, 2, rec.Attempts) recent, err := store.Recent(ctx, 10) require.NoError(t, err) require.Len(t, recent, 1, "同一期只能有一行台账") // 不同周期应各自一行 require.NoError(t, store.Save(ctx, &SecurityReportLog{ PeriodKey: "month:2026-10", Kind: "month", WindowStart: start, WindowEnd: end, Sent: true, })) recent, err = store.Recent(ctx, 10) require.NoError(t, err) require.Len(t, recent, 2) }