package trafficusage import ( "context" "database/sql" "database/sql/driver" "errors" "io" "strings" "testing" "time" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gcfg" "git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/library/tdengine" "git.5iot.com/iotlcm-public/iotlcm-admin-svr/internal/model" ) type routingDriver struct{ statements []string } type routingConnection struct{ owner *routingDriver } type routingRows struct { count bool emitted bool } func (recorder *routingDriver) Open(string) (driver.Conn, error) { return &routingConnection{owner: recorder}, nil } func (connection *routingConnection) Prepare(string) (driver.Stmt, error) { return nil, errors.New("unexpected prepare") } func (connection *routingConnection) Close() error { return nil } func (connection *routingConnection) Begin() (driver.Tx, error) { return nil, errors.New("unexpected transaction") } func (connection *routingConnection) Ping(context.Context) error { return nil } func (connection *routingConnection) ExecContext(_ context.Context, statement string, _ []driver.NamedValue) (driver.Result, error) { connection.owner.statements = append(connection.owner.statements, statement) return driver.RowsAffected(1), nil } func (connection *routingConnection) QueryContext(_ context.Context, statement string, _ []driver.NamedValue) (driver.Rows, error) { connection.owner.statements = append(connection.owner.statements, statement) return &routingRows{count: strings.Contains(statement, "COUNT(*)")}, nil } func (rows *routingRows) Columns() []string { if rows.count { return []string{"count"} } return []string{"ts", "usage_mb", "package_code", "iccid"} } func (rows *routingRows) Close() error { return nil } func (rows *routingRows) Next(values []driver.Value) error { if !rows.count || rows.emitted { return io.EOF } rows.emitted = true values[0] = int64(0) return nil } var routingRecorder = &routingDriver{} func init() { sql.Register("admin_traffic_split_routing", routingRecorder) } func TestTrafficUsageReadWriteDeleteSplitRouting(t *testing.T) { recorder := routingRecorder recorder.statements = nil previous := g.Cfg().GetAdapter() t.Cleanup(func() { _ = tdengine.Close(); g.Cfg().SetAdapter(previous) }) adapter, err := gcfg.NewAdapterContent("tsdb:\n database: TdEngine\n tdengine:\n type: admin_traffic_split_routing\n dsn: test\n dbName: legacy\n dailyDbName: iotlcm_daily\n monthlyDbName: iotlcm_monthly\n") if err != nil { t.Fatal(err) } g.Cfg().SetAdapter(adapter) _ = tdengine.Close() usage := NewTrafficUsage() ctx := context.Background() period := time.Date(2026, 10, 21, 0, 0, 0, 0, time.FixedZone("CST", 8*60*60)) base := model.TrafficUsageBase{Iccid: "8986112520609988658", UsageMb: 12.5, PackageCode: "test"} for _, test := range []struct { name string database string stable string operation func() error }{ {"write_daily", "iotlcm_daily", "sim_card_usage_d_", func() error { return usage.WriteDaily(ctx, []model.TrafficUsageDailyItem{{TrafficUsageBase: base, Date: period}}) }}, {"write_monthly", "iotlcm_monthly", "sim_card_usage_m_", func() error { return usage.WriteMonthly(ctx, []model.TrafficUsageMonthlyItem{{TrafficUsageBase: base, Month: period}}) }}, {"read_daily", "iotlcm_daily", "sim_card_usage_d_", func() error { _, _, err := usage.QueryDaily(ctx, model.TrafficUsageQueryInput{Iccid: base.Iccid}) return err }}, {"read_monthly", "iotlcm_monthly", "sim_card_usage_m_", func() error { _, _, err := usage.QueryMonthly(ctx, model.TrafficUsageQueryInput{Iccid: base.Iccid}) return err }}, {"delete_daily", "iotlcm_daily", "sim_card_usage_d_", func() error { return usage.DeleteDaily(ctx, model.TrafficUsageDeleteInput{Iccid: base.Iccid, StartTime: period, EndTime: period.AddDate(0, 0, 1)}) }}, {"delete_monthly", "iotlcm_monthly", "sim_card_usage_m_", func() error { return usage.DeleteMonthly(ctx, model.TrafficUsageDeleteInput{Iccid: base.Iccid, StartTime: period, EndTime: period.AddDate(0, 1, 0)}) }}, } { t.Run(test.name, func(t *testing.T) { recorder.statements = nil if err := test.operation(); err != nil { t.Fatal(err) } if len(recorder.statements) == 0 { t.Fatal("no SQL issued") } for _, statement := range recorder.statements { normalized := strings.ReplaceAll(statement, "\x60", "") if !strings.Contains(normalized, test.database+"."+test.stable+base.Iccid) || strings.Contains(normalized, "legacy.") { t.Errorf("wrong target database: %s", statement) } } }) } }