package trafficusage import ( "context" "database/sql" "database/sql/driver" "errors" "strings" "testing" "time" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gcfg" "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/library/tdengine" "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model" ) type routingDriver struct{ statements []string } type routingConnection struct{ owner *routingDriver } 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 } var routingRecorder = &routingDriver{} func init() { sql.Register("traffic_usage_split_routing", routingRecorder) } func TestRepositoryWritesToConfiguredSplitDatabase(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: traffic_usage_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() repository := New() 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"} if err := repository.UpdateDaily(context.Background(), []model.TrafficUsageDailyItem{{TrafficUsageBase: base, Date: period}}); err != nil { t.Fatal(err) } if err := repository.UpdateMonthly(context.Background(), []model.TrafficUsageMonthlyItem{{TrafficUsageBase: base, Month: period}}); err != nil { t.Fatal(err) } if len(recorder.statements) != 2 { t.Fatalf("statements=%v", recorder.statements) } for index, expected := range []string{ "INSERT INTO iotlcm_daily.sim_card_usage_d_8986112520609988658 USING iotlcm_daily.t_sim_card_usage_daily", "INSERT INTO iotlcm_monthly.sim_card_usage_m_8986112520609988658 USING iotlcm_monthly.t_sim_card_usage_monthly", } { if !strings.HasPrefix(recorder.statements[index], expected) { t.Errorf("statement=%s want prefix=%s", recorder.statements[index], expected) } } if err := repository.UpdateDaily(context.Background(), nil); err != nil { t.Fatal(err) } if err := repository.UpdateMonthly(context.Background(), nil); err != nil { t.Fatal(err) } if len(recorder.statements) != 2 { t.Fatal("empty updates issued SQL") } }