package cassandra import ( "context" "fmt" "log" "os" "sync/atomic" "testing" "time" "github.com/gocql/gocql" "github.com/testcontainers/testcontainers-go" "github.com/testcontainers/testcontainers-go/wait" ) type Container struct { Ctx context.Context Container testcontainers.Container Host string Port int } var cassandraDBTest *CassandraDB var keyspaceSequence atomic.Int64 func TestMain(m *testing.M) { container, db := connCassandraForTest() cassandraDBTest = db code := m.Run() cassandraDBTest.Close() if err := container.Container.Terminate(container.Ctx); err != nil { log.Fatalf("Failed to terminate Cassandra container: %v", err) } log.Println("[TEST] Container terminated") os.Exit(code) } func initCassandraContainer(version string) (Container, error) { ctx := context.Background() req := testcontainers.ContainerRequest{ Image: fmt.Sprintf("cassandra:%s", version), Env: map[string]string{ "CASSANDRA_START_RPC": "true", "CASSANDRA_NUM_TOKENS": "1", "CASSANDRA_ENDPOINT_SNITCH": "GossipingPropertyFileSnitch", "CASSANDRA_DC": "datacenter1", "CASSANDRA_RACK": "rack1", "MAX_HEAP_SIZE": "256M", "HEAP_NEWSIZE": "100M", }, ExposedPorts: []string{"9042/tcp"}, // 等待 Cassandra 啟動完成的指標字串,依據實際啟動 log 可調整 WaitingFor: wait.ForLog("Created default superuser role 'cassandra'"). WithStartupTimeout(2 * time.Minute), } cassandraContainer, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ ContainerRequest: req, Started: true, }) if err != nil { return Container{}, err } host, err := cassandraContainer.Host(ctx) if err != nil { return Container{}, err } mappedPort, err := cassandraContainer.MappedPort(ctx, "9042") if err != nil { return Container{}, err } return Container{ctx, cassandraContainer, host, mappedPort.Int()}, nil } func connCassandraForTest() (Container, *CassandraDB) { // 啟動 Cassandra container dbContainer, err := initCassandraContainer("5.0.4") if err != nil { log.Fatalf("Failed to initialize Cassandra container: %v", err) } db, err := NewCassandraDB( []string{dbContainer.Host}, WithPort(dbContainer.Port), WithConsistency(gocql.One), WithNumConns(5), ) if err != nil { log.Fatalf("Failed to initialize Cassandra DB: %v", err) } // 建立 keyspace 和 table err = db.EnsureTable(` CREATE KEYSPACE IF NOT EXISTS my_keyspace WITH replication = { 'class': 'SimpleStrategy', 'replication_factor': 1 };`) if err != nil { log.Fatalf("Failed to create keyspace: %v", err) } err = db.EnsureTable(` CREATE TABLE IF NOT EXISTS my_keyspace.monkey_entity ( id UUID, name TEXT, update_at TIMESTAMP, create_at TIMESTAMP, PRIMARY KEY ((id), name) );`) if err != nil { log.Fatalf("Failed to create table: %v", err) } return dbContainer, db } func generateRandomKeySpace(t *testing.T) string { ks := fmt.Sprintf("my_keyspace_%d", keyspaceSequence.Add(1)) err := cassandraDBTest.EnsureTable(fmt.Sprintf(` CREATE KEYSPACE IF NOT EXISTS %s WITH replication = { 'class': 'SimpleStrategy', 'replication_factor': 1 };`, ks)) if err != nil { t.Fatalf("Failed to create keyspace: %v", err) } err = cassandraDBTest.EnsureTable(fmt.Sprintf(` CREATE TABLE IF NOT EXISTS %s.monkey_entity ( id UUID, name TEXT, update_at TIMESTAMP, create_at TIMESTAMP, PRIMARY KEY ((id), name) );`, ks)) if err != nil { log.Fatalf("Failed to create table: %v", err) } return ks } // Animal 為不實作 TableName 方法的範例 struct,則會以型別名稱轉換成 snake_case type Animal struct { ID gocql.UUID `db:"id" partition_key:"true"` Type string `db:"type"` } func (m *Animal) TableName() string { return "animal" } // InvalidEntity 為無 partition key 的範例 struct,預期產生錯誤 type InvalidEntity struct { Field string `db:"field"` } type MonkeyEntity struct { ID gocql.UUID `db:"id" partition_key:"true"` Name string `db:"name" clustering_key:"true" sai:"true"` UpdateAt time.Time `db:"update_at"` CreateAt time.Time `db:"create_at"` } func (m *MonkeyEntity) TableName() string { return "monkey_entity" } type CatEntity struct { ID *gocql.UUID `db:"id" partition_key:"true"` Name *string `db:"name" partition_key:"true"` UpdateAt *time.Time `db:"update_at"` CreateAt *time.Time `db:"create_at" clustering_key:"true"` } func (m *CatEntity) TableName() string { return "cat_entity" } type Consistency struct { ID gocql.UUID `db:"id" partition_key:"true"` ConsistencyName string `db:"consistency_name"` // can editor ConsistencyType string `db:"consistency_type"` LastTaskID string `db:"last_task_id"` // ConsistencyTask ID Target string `db:"target"` // file name can editor Status string `db:"status"` ConsistencyMap string `db:"consistency_map"` // JSON string CreateAT int64 `db:"create_at"` UpdateAT int64 `db:"update_at"` } func (c *Consistency) TableName() string { return "consistency" }