backend/pkg/library/cassandra/cassandra_test.go

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"
}