wx_service/internal/database/smoke_index_migration_test.go

package database

import (
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func newMockDB(t *testing.T) (*gorm.DB, sqlmock.Sqlmock, func()) {
	t.Helper()

	sqlDB, mock, err := sqlmock.New()
	if err != nil {
		t.Fatalf("sqlmock.New: %v", err)
	}

	gdb, err := gorm.Open(mysql.New(mysql.Config{
		Conn:                      sqlDB,
		SkipInitializeWithVersion: true,
	}), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	if err != nil {
		_ = sqlDB.Close()
		t.Fatalf("gorm.Open: %v", err)
	}

	return gdb, mock, func() { _ = sqlDB.Close() }
}

func TestRepairSmokeAINextSmokeIndexesRecreatesBrokenIndex(t *testing.T) {
	t.Parallel()

	db, mock, cleanup := newMockDB(t)
	defer cleanup()

	mock.ExpectQuery("SHOW INDEX FROM `fa_smoke_ai_next_smoke` WHERE Key_name = \\?").
		WithArgs("uniq_smoke_ai_next_node").
		WillReturnRows(
			sqlmock.NewRows([]string{"Key_name", "Seq_in_index", "Column_name"}).
				AddRow("uniq_smoke_ai_next_node", 1, "node_type").
				AddRow("uniq_smoke_ai_next_node", 2, "node_at"),
		)
	mock.ExpectExec("ALTER TABLE `fa_smoke_ai_next_smoke` DROP INDEX `uniq_smoke_ai_next_node`").
		WillReturnResult(sqlmock.NewResult(0, 0))
	mock.ExpectExec("ALTER TABLE `fa_smoke_ai_next_smoke` ADD UNIQUE KEY `uniq_smoke_ai_next_node` \\(`ai_advice_id`,`node_type`,`node_at`\\)").
		WillReturnResult(sqlmock.NewResult(0, 0))

	if err := repairSmokeAINextSmokeIndexes(db); err != nil {
		t.Fatalf("repairSmokeAINextSmokeIndexes: %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("unmet expectations: %v", err)
	}
}

func TestRepairSmokeAINextSmokeIndexesKeepsCorrectIndex(t *testing.T) {
	t.Parallel()

	db, mock, cleanup := newMockDB(t)
	defer cleanup()

	mock.ExpectQuery("SHOW INDEX FROM `fa_smoke_ai_next_smoke` WHERE Key_name = \\?").
		WithArgs("uniq_smoke_ai_next_node").
		WillReturnRows(
			sqlmock.NewRows([]string{"Key_name", "Seq_in_index", "Column_name"}).
				AddRow("uniq_smoke_ai_next_node", 1, "ai_advice_id").
				AddRow("uniq_smoke_ai_next_node", 2, "node_type").
				AddRow("uniq_smoke_ai_next_node", 3, "node_at"),
		)

	if err := repairSmokeAINextSmokeIndexes(db); err != nil {
		t.Fatalf("repairSmokeAINextSmokeIndexes: %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("unmet expectations: %v", err)
	}
}

func TestRepairSmokeAIAdviceIndexesRecreatesBrokenIndex(t *testing.T) {
	t.Parallel()

	db, mock, cleanup := newMockDB(t)
	defer cleanup()

	mock.ExpectQuery("SHOW INDEX FROM `fa_smoke_ai_advice` WHERE Key_name = \\?").
		WithArgs("uniq_smoke_ai_advice").
		WillReturnRows(
			sqlmock.NewRows([]string{"Key_name", "Seq_in_index", "Column_name"}).
				AddRow("uniq_smoke_ai_advice", 1, "uid").
				AddRow("uniq_smoke_ai_advice", 2, "advice_date").
				AddRow("uniq_smoke_ai_advice", 3, "prompt_version"),
		)
	mock.ExpectExec("ALTER TABLE `fa_smoke_ai_advice` DROP INDEX `uniq_smoke_ai_advice`").
		WillReturnResult(sqlmock.NewResult(0, 0))
	mock.ExpectExec("ALTER TABLE `fa_smoke_ai_advice` ADD UNIQUE KEY `uniq_smoke_ai_advice` \\(`uid`,`type`,`advice_date`,`prompt_version`\\)").
		WillReturnResult(sqlmock.NewResult(0, 0))

	if err := repairSmokeAIAdviceIndexes(db); err != nil {
		t.Fatalf("repairSmokeAIAdviceIndexes: %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("unmet expectations: %v", err)
	}
}

func TestRepairSmokeAIAdviceIndexesKeepsCorrectIndex(t *testing.T) {
	t.Parallel()

	db, mock, cleanup := newMockDB(t)
	defer cleanup()

	mock.ExpectQuery("SHOW INDEX FROM `fa_smoke_ai_advice` WHERE Key_name = \\?").
		WithArgs("uniq_smoke_ai_advice").
		WillReturnRows(
			sqlmock.NewRows([]string{"Key_name", "Seq_in_index", "Column_name"}).
				AddRow("uniq_smoke_ai_advice", 1, "uid").
				AddRow("uniq_smoke_ai_advice", 2, "type").
				AddRow("uniq_smoke_ai_advice", 3, "advice_date").
				AddRow("uniq_smoke_ai_advice", 4, "prompt_version"),
		)

	if err := repairSmokeAIAdviceIndexes(db); err != nil {
		t.Fatalf("repairSmokeAIAdviceIndexes: %v", err)
	}
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Fatalf("unmet expectations: %v", err)
	}
}