usage_log_repo_unit_test.go 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. //go:build unit
  2. package repository
  3. import (
  4. "strings"
  5. "testing"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/internal/service"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestSafeDateFormat(t *testing.T) {
  11. tests := []struct {
  12. name string
  13. granularity string
  14. expected string
  15. }{
  16. // 合法值
  17. {"hour", "hour", "YYYY-MM-DD HH24:00"},
  18. {"day", "day", "YYYY-MM-DD"},
  19. {"week", "week", "IYYY-IW"},
  20. {"month", "month", "YYYY-MM"},
  21. // 非法值回退到默认
  22. {"空字符串", "", "YYYY-MM-DD"},
  23. {"未知粒度 year", "year", "YYYY-MM-DD"},
  24. {"未知粒度 minute", "minute", "YYYY-MM-DD"},
  25. // 恶意字符串
  26. {"SQL 注入尝试", "'; DROP TABLE users; --", "YYYY-MM-DD"},
  27. {"带引号", "day'", "YYYY-MM-DD"},
  28. {"带括号", "day)", "YYYY-MM-DD"},
  29. {"Unicode", "日", "YYYY-MM-DD"},
  30. }
  31. for _, tc := range tests {
  32. t.Run(tc.name, func(t *testing.T) {
  33. got := safeDateFormat(tc.granularity)
  34. require.Equal(t, tc.expected, got, "safeDateFormat(%q)", tc.granularity)
  35. })
  36. }
  37. }
  38. func TestBuildUsageLogBatchInsertQuery_UsesConflictDoNothing(t *testing.T) {
  39. log := &service.UsageLog{
  40. UserID: 1,
  41. APIKeyID: 2,
  42. AccountID: 3,
  43. RequestID: "req-batch-no-update",
  44. Model: "gpt-5",
  45. InputTokens: 10,
  46. OutputTokens: 5,
  47. TotalCost: 1.2,
  48. ActualCost: 1.2,
  49. CreatedAt: time.Now().UTC(),
  50. }
  51. prepared := prepareUsageLogInsert(log)
  52. query, _ := buildUsageLogBatchInsertQuery([]string{usageLogBatchKey(log.RequestID, log.APIKeyID)}, map[string]usageLogInsertPrepared{
  53. usageLogBatchKey(log.RequestID, log.APIKeyID): prepared,
  54. })
  55. require.Contains(t, query, "ON CONFLICT (request_id, api_key_id) DO NOTHING")
  56. require.NotContains(t, strings.ToUpper(query), "DO UPDATE")
  57. }