package sqltx

import (
	"regexp"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
)

func TestConstraintCheckingPessimisticTxRestoresPriorSessionValue(t *testing.T) {
	for _, previous := range []string{"OFF", "ON"} {
		t.Run(previous, func(t *testing.T) {
			db, mock, err := sqlmock.New()
			if err != nil {
				t.Fatalf("sqlmock.New(): %v", err)
			}
			defer db.Close()

			mock.ExpectQuery(regexp.QuoteMeta(tidbConstraintCheckInPlaceSelect)).
				WillReturnRows(sqlmock.NewRows([]string{"constraint_check"}).AddRow(previous))
			mock.ExpectExec(regexp.QuoteMeta("SET SESSION tidb_txn_mode = 'pessimistic'")).
				WillReturnResult(sqlmock.NewResult(0, 0))
			mock.ExpectExec(regexp.QuoteMeta(tidbConstraintCheckInPlaceOn)).
				WillReturnResult(sqlmock.NewResult(0, 0))
			mock.ExpectBegin()
			mock.ExpectRollback()
			mock.ExpectExec(regexp.QuoteMeta("SET SESSION tidb_constraint_check_in_place_pessimistic = " + previous)).
				WillReturnResult(sqlmock.NewResult(0, 0))
			mock.ExpectExec(regexp.QuoteMeta("SET SESSION tidb_txn_mode = DEFAULT")).
				WillReturnResult(sqlmock.NewResult(0, 0))

			tx, cleanup, err := BeginPessimisticConstraintCheckingTx(t.Context(), db)
			if err != nil {
				t.Fatalf("BeginPessimisticConstraintCheckingTx(): %v", err)
			}
			if err := tx.Rollback(); err != nil {
				t.Fatalf("Rollback(): %v", err)
			}
			if err := cleanup(); err != nil {
				t.Fatalf("cleanup(): %v", err)
			}
			if err := mock.ExpectationsWereMet(); err != nil {
				t.Fatalf("sql expectations: %v", err)
			}
		})
	}
}