package tests import ( "encoding/json" "net/http" "testing" "time" "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/pingcap/dbaas/billing2/api" "github.com/pingcap/dbaas/billing2/db" "github.com/pingcap/dbaas/billing2/db/entity" "github.com/pingcap/dbaas/billing2/forms" "github.com/pingcap/dbaas/billing2/handler/awshandler" "github.com/pingcap/dbaas/billing2/handler/centralhandler" "github.com/pingcap/dbaas/billing2/model/billingconfig" "github.com/pingcap/dbaas/billing2/model/marketing" "github.com/pingcap/dbaas/billing2/utils/ginutils" "github.com/pingcap/dbaas/billing2/utils/panicIf" "github.com/pingcap/dbaas/billing2/utils/tests" "github.com/pingcap/dbaas/billing2/utils/tests/datainit" "github.com/pingcap/dbaas/billing2/utils/timeutils" "github.com/shopspring/decimal" "github.com/stretchr/testify/suite" "github.com/stripe/stripe-go/v76" ) type ARTableTestSuite struct { tests.BillingTestSuite router *gin.Engine } func (suite *ARTableTestSuite) SetupTest() { suite.BillingTestSuite.SetupTest() suite.router = gin.New() suite.router.Use(ginutils.ErrorMiddleWare) api.InitializeARTableRouter(ginutils.MyRoute{Route: suite.router}) datainit.CreateDefaultUnits() datainit.CreateCategoryLabels() } func TestARTableTestSuite(t *testing.T) { suite.Run(t, new(ARTableTestSuite)) } // arTableMockStripeMonthlyInvoices, when set, replaces the fixed invoice list // returned by MockStripeHandler.GetMonthlyCreatedInvoices. AR generation tests // use it to simulate auto-pushed Stripe invoices that only carry merged_months // metadata (no invoice_start/invoice_end). Tests must reset it to nil. var arTableMockStripeMonthlyInvoices []*forms.StripeInvoice func (*MockStripeHandler) GetMonthlyCreatedInvoices(date time.Time) ([]*forms.StripeInvoice, error) { if arTableMockStripeMonthlyInvoices != nil { return arTableMockStripeMonthlyInvoices, nil } return []*forms.StripeInvoice{ { CustomerID: "123", Status: "PAID", Number: "123", SubTotal: decimal.NewFromInt(2), TaxRate: decimal.NewFromInt(10), Total: decimal.RequireFromString("2.2"), InvoicePDF: "0000.pdf", Paid: true, StartDate: datainit.BillingDate, Created: datainit.BillingDate.AddDate(0, 0, 2).Unix(), StatusTransitions: &stripe.InvoiceStatusTransitions{ FinalizedAt: datainit.BillingDate.AddDate(0, 0, 3).Unix(), PaidAt: datainit.BillingDate.AddDate(0, 0, 3).Unix(), }, }, { CustomerID: "123", Status: "PAID", Number: "124", SubTotal: decimal.NewFromInt(1), TaxRate: decimal.NewFromInt(10), Total: decimal.RequireFromString("1.1"), InvoicePDF: "0001.pdf", Paid: false, StartDate: datainit.BillingDate, Created: datainit.BillingDate.AddDate(0, 0, 8).Unix(), DueDate: 0, StatusTransitions: &stripe.InvoiceStatusTransitions{ FinalizedAt: datainit.BillingDate.AddDate(0, 0, 8).Unix(), }, }, }, nil } func (suite *ARTableTestSuite) TestInitARTable() { MockPaymentHandler() central := &MockCentralHandlerImpl{} tenants := []forms.Tenant{ { TenantID: 1, Kind: forms.TenantKindExternal, }, { TenantID: 2, Kind: forms.TenantKindExternal, }, { TenantID: 3, Kind: forms.TenantKindExternal, }, { TenantID: 4, Kind: forms.TenantKindInternal, }, { TenantID: 5, Kind: forms.TenantKindExternal, }, { TenantID: 6, Kind: forms.TenantKindExternal, }, { TenantID: 7, Kind: forms.TenantKindExternal, }, { TenantID: 8, Kind: forms.TenantKindExternal, }, { TenantID: 9, Kind: forms.TenantKindExternal, }, { TenantID: 10, Kind: forms.TenantKindExternal, }, { TenantID: 11, Kind: forms.TenantKindExternal, }, } central.SetAllTenants(tenants) centralhandler.SetCentralHandler(central) var err error assert := suite.Assert() timeutils.SetMockNow(datainit.BillingDate.AddDate(0, 1, 0)) billingConfigs := make([]*forms.BillingConfig, len(tenants)) for i, t := range tenants { switch t.TenantID { case 7, 9: billingConfigs[i], err = datainit.CreateBillingConfig(t.TenantID, entity.BillingConfigType_MarketplaceReseller) panicIf.Error(err) panicIf.Error(entity.NewBillingConfigQuerySet(db.GetDB()).IDEq(billingConfigs[i].BillingConfigID).GetUpdater().SetManagedTenantID(t.TenantID + 1).Update()) case 8, 10: billingConfigs[i], err = datainit.CreateBillingConfig(t.TenantID, entity.BillingConfigType_MarketplaceResellerManagedCustomer) panicIf.Error(err) default: billingConfigs[i], err = datainit.CreateBillingConfig(t.TenantID, entity.BillingConfigType_DirectCustomer) panicIf.Error(err) } } prepayTenantID := uint(2) noInvoiceTenantID := uint(5) billingConfigID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(prepayTenantID) assert.NoError(err) req := &forms.CreditsPackage{ StartDate: datainit.BillingDate.AddDate(0, 0, -15), ExpiredDate: datainit.BillingDate.AddDate(0, 1, 22), TotalCredits: decimal.NewFromInt(100), Usage: entity.CreditsPackageUsage_Prepayment, BindedBillingConfigID: &billingConfigID, } prepayCredits, err := (&marketing.CreditsPackageModel{}).CreateCreditsPackage(prepayTenantID, req) assert.NoError(err) targetUsages := []string{entity.CreditsPackageUsage_CustomerSuccess, entity.CreditsPackageUsage_Redeem, entity.CreditsPackageUsage_Program} for i, usage := range targetUsages { req := &forms.CreditsPackage{ StartDate: datainit.BillingDate.AddDate(0, 0, -15), ExpiredDate: datainit.BillingDate.AddDate(0, 1, 22), TotalCredits: decimal.NewFromInt(100), Usage: usage, BindedBillingConfigID: &billingConfigID, } usageCredits, err := (&marketing.CreditsPackageModel{}).CreateCreditsPackage(prepayTenantID, req) assert.NoError(err) basePrice := int64(1 + i) premiumCost := entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(basePrice * 5), RawPrice: decimal.NewFromInt(basePrice * 5), TenantID: prepayTenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *usageCredits.BindedPackageID, Unit: datainit.AWSPremiumBackupUnit, } assert.NoError(premiumCost.Create(db.GetDB())) for _, cost := range []entity.UserUnitCost{{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(basePrice), RawPrice: decimal.NewFromInt(basePrice), TenantID: prepayTenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *usageCredits.BindedPackageID, Unit: datainit.SupportCommitmentUnit, }, { CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(basePrice * 10), RawPrice: decimal.NewFromInt(basePrice * 10), TenantID: prepayTenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *usageCredits.BindedPackageID, Unit: datainit.AWSNodeUnit, }, { CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(basePrice * 2), RawPrice: decimal.NewFromInt(basePrice * 2), TenantID: prepayTenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *usageCredits.BindedPackageID, Unit: datainit.AWSServerlessStorageUnit, }, { CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(basePrice * 3), RawPrice: decimal.NewFromInt(basePrice * 3), TenantID: prepayTenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *usageCredits.BindedPackageID, Unit: datainit.AWSStarterRuUnit, }, { CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(11), RawPrice: decimal.NewFromInt(11), TenantID: prepayTenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *usageCredits.BindedPackageID, Unit: datainit.EssentialColumnStorageUnit, }} { err = cost.Create(db.GetDB()) assert.NoError(err) } } payments := []string{entity.PaymentType_STRIPE, entity.PaymentType_AWS, entity.PaymentType_WireTransfer, entity.PaymentType_WireTransfer, entity.PaymentType_WireTransfer, entity.PaymentType_MSP, entity.PaymentType_AWS, entity.PaymentType_AWSCPPOTransfer, entity.PaymentType_AWSCPPOReseller, entity.PaymentType_AWSCPPOTransfer, entity.PaymentType_AliCloud} assert.Len(payments, len(tenants)) for index, tenant := range tenants { paymentType := payments[index] tenantID := tenant.TenantID billingConfigID := billingConfigs[index].BillingConfigID p := &entity.PaymentMethod{ PaymentType: paymentType, Payable: true, PaymentCode: "123", TenantID: tenantID, Status: entity.PaymentMethodStatus_Active, BindedBillingConfigID: billingConfigID, Meta: &entity.PaymentMeta{ EnlargeMarketplaceMetering: billingConfigs[index].Type == entity.BillingConfigType_MarketplaceReseller, }, } if paymentType == entity.PaymentType_AWSCPPOReseller { subCode := "233333" p.PaymentSubCode = &subCode } err = p.Create(db.GetDB()) assert.NoError(err) if tenantID == noInvoiceTenantID { continue } supportCost := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(1), RawPrice: decimal.NewFromInt(1), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Invoice, BindedBillingConfigID: billingConfigID, Unit: datainit.SupportCommitmentUnit, } cost := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(1), RawPrice: decimal.NewFromInt(1), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Invoice, BindedBillingConfigID: billingConfigID, Unit: datainit.AWSNodeUnit, } premiumInvoiceCost := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(4), RawPrice: decimal.NewFromInt(4), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Invoice, BindedBillingConfigID: billingConfigID, Unit: datainit.AWSPremiumBackupUnit, } if paymentType == entity.PaymentType_AWS { cost2 := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.RequireFromString("1.7"), RawPrice: decimal.NewFromInt(2), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Invoice, BindedBillingConfigID: billingConfigID, Unit: datainit.AWSNodeUnit, } inv := &entity.Invoice{ InvoiceDate: datainit.BillingDate, CostPrice: cost2.CostPrice, RawPrice: cost2.RawPrice, Type: entity.InvoiceType_INMONTH, BindedBillingConfigID: billingConfigID, Status: entity.InvoiceStatus_SENT, TenantID: tenantID, ChannelLoss: decimal.RequireFromString("0.7"), } err = inv.Create(db.GetDB()) assert.NoError(err) cost2.InvoiceID = int(inv.ID) err = cost2.Create(db.GetDB()) assert.NoError(err) } if paymentType == entity.PaymentType_AliCloud { cost.Unit = datainit.EssentialColumnStorageUnit } if tenantID != prepayTenantID { invoice := &entity.Invoice{ InvoiceDate: datainit.BillingDate, CostPrice: decimal.NewFromInt(2), RawPrice: decimal.NewFromInt(2), Type: entity.InvoiceType_MONTHLY, BindedBillingConfigID: billingConfigID, Status: entity.InvoiceStatus_GENERATED, TenantID: tenantID, } if paymentType == entity.PaymentType_AliCloud { invoice.ChannelLoss = decimal.RequireFromString("0.3") } err = invoice.Create(db.GetDB()) assert.NoError(err) cost.InvoiceID = int(invoice.ID) supportCost.InvoiceID = int(invoice.ID) premiumInvoiceCost.InvoiceID = int(invoice.ID) } else { serverlessCost := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(1), RawPrice: decimal.NewFromInt(1), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *prepayCredits.BindedPackageID, Unit: datainit.AWSServerlessRUUnit, } err = serverlessCost.Create(db.GetDB()) assert.NoError(err) starterCost := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.NewFromInt(1), RawPrice: decimal.NewFromInt(1), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Credits, BindedBillingConfigID: billingConfigID, BindedPackageID: *prepayCredits.BindedPackageID, Unit: datainit.AWSStarterRuUnit, } err = starterCost.Create(db.GetDB()) assert.NoError(err) supportCost.ConsumedSource = entity.ConsumedSource_Credits supportCost.BindedPackageID = *prepayCredits.BindedPackageID cost.ConsumedSource = entity.ConsumedSource_Credits cost.BindedPackageID = *prepayCredits.BindedPackageID premiumInvoiceCost.ConsumedSource = entity.ConsumedSource_Credits premiumInvoiceCost.BindedPackageID = *prepayCredits.BindedPackageID } if paymentType == entity.PaymentType_AWSCPPOReseller { cost2 := &entity.UserUnitCost{ CostID: uuid.NewString(), CostTimeStamp: datainit.BillingDate.Unix(), CostPrice: decimal.RequireFromString("1.7"), RawPrice: decimal.NewFromInt(2), TenantID: tenantID, Status: entity.UserUnitCostStatus_DONE, ConsumedSource: entity.ConsumedSource_Invoice, BindedBillingConfigID: billingConfigID, Unit: datainit.AWSNodeUnit, } inv := &entity.Invoice{ InvoiceDate: datainit.BillingDate, CostPrice: cost2.CostPrice, RawPrice: cost2.RawPrice, Type: entity.InvoiceType_INMONTH, BindedBillingConfigID: billingConfigID, Status: entity.InvoiceStatus_SPLIT, TenantID: tenantID, } err = inv.Create(db.GetDB()) assert.NoError(err) cost2.InvoiceID = int(inv.ID) err = cost2.Create(db.GetDB()) assert.NoError(err) subInv := &entity.Invoice{ InvoiceDate: datainit.BillingDate, CostPrice: cost2.CostPrice, RawPrice: cost2.RawPrice, Type: entity.InvoiceType_SUBORDINATED, BindedBillingConfigID: billingConfigID, Status: entity.InvoiceStatus_SENT, TenantID: tenantID, PaymentCode: p.PaymentSubCode, ChannelLoss: decimal.RequireFromString("0.7"), } assert.NoError(subInv.Create(db.GetDB())) subInv2 := &entity.Invoice{ InvoiceDate: datainit.BillingDate, CostPrice: decimal.NewFromInt(2), RawPrice: decimal.NewFromInt(2), Type: entity.InvoiceType_SUBORDINATED, BindedBillingConfigID: billingConfigID, PaymentCode: &p.PaymentCode, Status: entity.InvoiceStatus_GENERATED, TenantID: tenantID, ChannelLoss: decimal.RequireFromString("0.53"), } assert.NoError(subInv2.Create(db.GetDB())) } err = cost.Create(db.GetDB()) assert.NoError(err) err = supportCost.Create(db.GetDB()) assert.NoError(err) err = premiumInvoiceCost.Create(db.GetDB()) assert.NoError(err) } w := tests.PerformRequest(suite.router, "GET", "/cr_tables/init?year=2021&month=11&day=1", nil) assert.Equal(http.StatusOK, w.Code, w.Body.String()) expected := `bill_month,start_date,end_date,tenant_id,stripe_amount,stripe_tax_rate,stripe_amount_including_tax,stripe_invoice_number,wire_transfer_amount,marketplace_amount,prepayment_credits,total_amount,invoice_amount,is_estimated,dedicated_accrual_amount,support_accrual_amount,stripe_invoice,stripe_invoice_paid,stripe_invoice_finalized_date,stripe_invoice_paid_date,raw_starter_amount,raw_dedicated_amount,raw_essential_amount,raw_support_amount,raw_premium_amount,raw_lake_amount,raw_byoc_amount,serverless_amount,dedicated_amount,starter_amount,essential_amount,support_amount,premium_amount,lake_amount,byoc_amount,raw_program_credit_amount_starter,raw_program_credit_amount_dedicated,raw_program_credit_amount_essential,raw_program_credit_amount_support,raw_program_credit_amount_premium,raw_program_credit_amount_lake,raw_program_credit_amount_byoc,raw_redeem_credit_amount_starter,raw_redeem_credit_amount_dedicated,raw_redeem_credit_amount_essential,raw_redeem_credit_amount_support,raw_redeem_credit_amount_premium,raw_redeem_credit_amount_lake,raw_redeem_credit_amount_byoc,raw_customer_success_credit_amount_starter,raw_customer_success_credit_amount_dedicated,raw_customer_success_credit_amount_essential,raw_customer_success_credit_amount_support,raw_customer_success_credit_amount_premium,raw_customer_success_credit_amount_lake,raw_customer_success_credit_amount_byoc,channel_loss_amount,starter_allocated_loss,dedicated_allocated_loss,essential_allocated_loss,support_allocated_loss,premium_allocated_loss,lake_allocated_loss,byoc_allocated_loss,cppo_discountable_product_amount,cppo_non_discountable_product_amount 2021-11,2021-11-01,2021-11-30,1,"[""2"",""1""]","[""10"",""10""]","[""2.2"",""1.1""]","[""123"",""124""]",,,,6,6,false,,,"[""0000.pdf"",""0001.pdf""]","[true,false]","[""2021-11-04"",""2021-11-09""]","[""2021-11-04"",""null""]",0,1,0,1,4,0,0,0,1,0,0,1,4,0,0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, 2021-11,2021-11-01,2021-11-30,2,,,,,,1.7,8,9,9.7,false,,,,,,,1,2.7,0,1,4,0,0,0.93,2.51,0.93,0,0.93,3.7,0,0,9,30,11,3,15,0,0,6,20,11,2,10,0,0,3,10,11,1,5,0,0,0.7,0.07,0.19,0,0.07,0.3,0,0,, 2021-11,2021-11-01,2021-11-30,3,,,,,6,,,6,6,false,,,,,,,0,1,0,1,4,0,0,0,1,0,0,1,4,0,0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, 2021-11,2021-11-01,2021-11-30,8,,,,,,7.7,,7,7.7,false,,,,,,,0,2.7,0,1,4,0,0,0,2.46,0,0,0.91,3.63,0,0,,,,,,,,,,,,,,,,,,,,,,0.7,0,0.24,0,0.09,0.37,0,0,, 2021-11,2021-11-01,2021-11-30,10,,,,,,7.7,,6.47,7.7,false,,,,,,,0,2.7,0,1,4,0,0,0,2.27,0,0,0.85,3.35,0,0,,,,,,,,,,,,,,,,,,,,,,1.23,0,0.43,0,0.15,0.65,0,0,1.47,1 2021-11,2021-11-01,2021-11-30,11,,,,,,6,,5.7,6,false,,,,,,,0,0,1,1,4,0,0,0,0,0,0.95,0.95,3.8,0,0,,,,,,,,,,,,,,,,,,,,,,0.3,0,0,0.05,0.05,0.2,0,0,, ` assert.Equal(expected, w.Body.String()) } func (suite *ARTableTestSuite) TestARTableExportStatusAndDownload() { assert := suite.Assert() now := time.Date(2026, 5, 29, 13, 0, 0, 0, time.UTC) date := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC) finalKey := awshandler.BuildS3ARTableKey(date, "actual") markerKey := awshandler.BuildS3ARTableRunningKey(date, "actual") markers := map[string]*awshandler.ARTableExportMarker{} lastModified := map[string]time.Time{} restore := api.SetARTableExportHooksForTest(&api.ARTableExportHooks{ Now: func() time.Time { return now }, GetMarker: func(key string) (*awshandler.ARTableExportMarker, bool, error) { marker, ok := markers[key] if !ok { return nil, false, nil } ret := *marker return &ret, true, nil }, PutMarker: func(key string, marker *awshandler.ARTableExportMarker) error { ret := *marker markers[key] = &ret return nil }, DeleteMarker: func(key string) error { delete(markers, key) return nil }, GetLastMod: func(key string) (*time.Time, bool, error) { t, ok := lastModified[key] if !ok { return nil, false, nil } ret := t return &ret, true, nil }, GetPresignURL: func(key string) (string, error) { return "https://example.com/download?key=" + key, nil }, StartExportJob: func(time.Time, string, string, string, *awshandler.ARTableExportMarker) {}, }) defer restore() w := tests.PerformGetRequest(suite.router, "/ar_tables/export/status?year=2026&month=5&type=actual") assert.Equal(http.StatusOK, w.Code, w.Body.String()) statusResp := &forms.ARTableExportStatusResponse{} assert.NoError(json.Unmarshal(w.Body.Bytes(), statusResp)) assert.Equal("not_found", statusResp.Status) assert.Equal(finalKey, statusResp.Key) markers[markerKey] = &awshandler.ARTableExportMarker{ RunID: "running", StartedAt: now.Add(-2 * time.Minute), TimeoutAt: now.Add(13 * time.Minute), Year: 2026, Month: 5, Type: "actual", } w = tests.PerformGetRequest(suite.router, "/ar_tables/export/status?year=2026&month=5&type=actual") assert.Equal(http.StatusOK, w.Code, w.Body.String()) assert.NoError(json.Unmarshal(w.Body.Bytes(), statusResp)) assert.Equal("running", statusResp.Status) lastModified[finalKey] = now w = tests.PerformGetRequest(suite.router, "/ar_tables/export/status?year=2026&month=5&type=actual") assert.Equal(http.StatusOK, w.Code, w.Body.String()) assert.NoError(json.Unmarshal(w.Body.Bytes(), statusResp)) assert.Equal("ready", statusResp.Status) delete(lastModified, finalKey) markers[markerKey] = &awshandler.ARTableExportMarker{ RunID: "timeout", StartedAt: now.Add(-30 * time.Minute), TimeoutAt: now.Add(-15 * time.Minute), Year: 2026, Month: 5, Type: "actual", } w = tests.PerformGetRequest(suite.router, "/ar_tables/export/status?year=2026&month=5&type=actual") assert.Equal(http.StatusOK, w.Code, w.Body.String()) assert.NoError(json.Unmarshal(w.Body.Bytes(), statusResp)) assert.Equal("timeout", statusResp.Status) delete(markers, markerKey) lastModified[finalKey] = now w = tests.PerformGetRequest(suite.router, "/ar_tables/export/download?year=2026&month=5&type=actual") assert.Equal(http.StatusOK, w.Code, w.Body.String()) downloadResp := &forms.ARTableExportDownloadResponse{} assert.NoError(json.Unmarshal(w.Body.Bytes(), downloadResp)) assert.Equal(finalKey, downloadResp.Key) assert.Equal("https://example.com/download?key="+finalKey, downloadResp.URL) markers[markerKey] = &awshandler.ARTableExportMarker{ RunID: "running-again", StartedAt: now.Add(-1 * time.Minute), TimeoutAt: now.Add(14 * time.Minute), Year: 2026, Month: 5, Type: "actual", } lastModified[finalKey] = now.Add(-2 * time.Minute) w = tests.PerformGetRequest(suite.router, "/ar_tables/export/download?year=2026&month=5&type=actual") assert.Equal(http.StatusBadRequest, w.Code, w.Body.String()) delete(markers, markerKey) delete(lastModified, finalKey) w = tests.PerformGetRequest(suite.router, "/ar_tables/export/download?year=2026&month=5&type=actual") assert.Equal(http.StatusNotFound, w.Code, w.Body.String()) } func (suite *ARTableTestSuite) TestARTableTriggerDeduplicatesRunningExport() { assert := suite.Assert() now := time.Date(2026, 5, 29, 13, 0, 0, 0, time.UTC) markers := map[string]*awshandler.ARTableExportMarker{} started := 0 restore := api.SetARTableExportHooksForTest(&api.ARTableExportHooks{ Now: func() time.Time { return now }, GetMarker: func(key string) (*awshandler.ARTableExportMarker, bool, error) { marker, ok := markers[key] if !ok { return nil, false, nil } ret := *marker return &ret, true, nil }, PutMarker: func(key string, marker *awshandler.ARTableExportMarker) error { ret := *marker markers[key] = &ret return nil }, DeleteMarker: func(key string) error { delete(markers, key) return nil }, GetLastMod: func(string) (*time.Time, bool, error) { return nil, false, nil }, GetPresignURL: func(key string) (string, error) { return "https://example.com/download?key=" + key, nil }, StartExportJob: func(time.Time, string, string, string, *awshandler.ARTableExportMarker) { started++ }, }) defer restore() w := tests.PerformRequest(suite.router, "POST", "/ar_tables/trigger?year=2026&month=5", nil) assert.Equal(http.StatusOK, w.Code, w.Body.String()) triggerResp := &forms.ARTableTriggerResponse{} assert.NoError(json.Unmarshal(w.Body.Bytes(), triggerResp)) assert.Equal("running", triggerResp.Status) assert.Equal(1, started) assert.Len(markers, 1) w = tests.PerformRequest(suite.router, "POST", "/ar_tables/trigger?year=2026&month=5", nil) assert.Equal(http.StatusOK, w.Code, w.Body.String()) assert.NoError(json.Unmarshal(w.Body.Bytes(), triggerResp)) assert.Equal("running", triggerResp.Status) assert.Equal(1, started) assert.Len(markers, 1) }