sync: 同步 v1.0.4 开源代码并脱敏

internal/repository/trafficusage/routing_test.go

@@ -0,0 +1,84 @@
+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")
+	}
+}