ops_ingress_reject_repo_test.go 1.2 KB

1234567891011121314151617181920212223242526272829303132
  1. package repository
  2. import (
  3. "context"
  4. "testing"
  5. "time"
  6. "github.com/DATA-DOG/go-sqlmock"
  7. "github.com/Wei-Shaw/sub2api/internal/service"
  8. "github.com/stretchr/testify/require"
  9. )
  10. func TestBatchUpsertIngressRejectsUsesFixedMultiRowChunks(t *testing.T) {
  11. db, mock, err := sqlmock.New()
  12. require.NoError(t, err)
  13. t.Cleanup(func() { _ = db.Close() })
  14. repo := &opsRepository{db: db}
  15. now := time.Now().UTC().Truncate(time.Minute)
  16. items := make([]*service.OpsIngressRejectAggregate, ingressRejectUpsertChunkSize+1)
  17. for i := range items {
  18. items[i] = &service.OpsIngressRejectAggregate{
  19. BucketStart: now, RejectReason: "invalid_api_key", RouteFamily: "messages",
  20. Protocol: "anthropic", ClientIP: "192.0.2.1", RequestCount: 1, FirstSeen: now, LastSeen: now,
  21. }
  22. }
  23. mock.ExpectBegin()
  24. mock.ExpectExec("INSERT INTO ops_ingress_reject_aggregates").WillReturnResult(sqlmock.NewResult(0, int64(ingressRejectUpsertChunkSize)))
  25. mock.ExpectExec("INSERT INTO ops_ingress_reject_aggregates").WillReturnResult(sqlmock.NewResult(0, 1))
  26. mock.ExpectCommit()
  27. require.NoError(t, repo.BatchUpsertIngressRejects(context.Background(), items))
  28. require.NoError(t, mock.ExpectationsWereMet())
  29. }