From 9de1261f3188d65c015e297cd67f1e768a696988 Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Tue, 15 Aug 2023 08:26:42 +0100
Subject: [PATCH] Add tests for the set unique operation

---
 pkg/migrations/op_common_test.go     |  7 ++
 pkg/migrations/op_set_unique_test.go | 96 ++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+)
 create mode 100644 pkg/migrations/op_set_unique_test.go

diff --git a/pkg/migrations/op_common_test.go b/pkg/migrations/op_common_test.go
index c81fdb92..8b474fe0 100644
--- a/pkg/migrations/op_common_test.go
+++ b/pkg/migrations/op_common_test.go
@@ -250,6 +250,13 @@ func ConstraintMustNotExist(t *testing.T, db *sql.DB, schema, table, constraint
 	}
 }
 
+func ConstraintMustExist(t *testing.T, db *sql.DB, schema, table, constraint string) {
+	t.Helper()
+	if !constraintExists(t, db, schema, table, constraint) {
+		t.Fatalf("Expected constraint %q to exist", constraint)
+	}
+}
+
 func IndexMustExist(t *testing.T, db *sql.DB, schema, table, index string) {
 	t.Helper()
 	if !indexExists(t, db, schema, table, index) {
diff --git a/pkg/migrations/op_set_unique_test.go b/pkg/migrations/op_set_unique_test.go
new file mode 100644
index 00000000..d2664b99
--- /dev/null
+++ b/pkg/migrations/op_set_unique_test.go
@@ -0,0 +1,96 @@
+package migrations_test
+
+import (
+	"database/sql"
+	"testing"
+
+	"pg-roll/pkg/migrations"
+)
+
+func TestSetColumnsUnique(t *testing.T) {
+	t.Parallel()
+
+	ExecuteTests(t, TestCases{{
+		name: "set unique",
+		migrations: []migrations.Migration{
+			{
+				Name: "01_add_table",
+				Operations: migrations.Operations{
+					&migrations.OpCreateTable{
+						Name: "reviews",
+						Columns: []migrations.Column{
+							{
+								Name:       "id",
+								Type:       "serial",
+								PrimaryKey: true,
+							},
+							{
+								Name:     "username",
+								Type:     "text",
+								Nullable: false,
+							},
+							{
+								Name:     "product",
+								Type:     "text",
+								Nullable: false,
+							},
+							{
+								Name:     "review",
+								Type:     "text",
+								Nullable: false,
+							},
+						},
+					},
+				},
+			},
+			{
+				Name: "02_set_unique",
+				Operations: migrations.Operations{
+					&migrations.OpSetUnique{
+						Table:   "reviews",
+						Columns: []string{"username", "product"},
+					},
+				},
+			},
+		},
+		afterStart: func(t *testing.T, db *sql.DB) {
+			// The unique index has been created on the underlying table.
+			idxName := migrations.IndexName("reviews", []string{"username", "product"})
+			IndexMustExist(t, db, "public", "reviews", idxName)
+
+			// Inserting values into the old schema that violate uniqueness should fail.
+			MustInsert(t, db, "public", "01_add_table", "reviews", map[string]string{
+				"username": "alice", "product": "apple", "review": "good",
+			})
+			MustNotInsert(t, db, "public", "01_add_table", "reviews", map[string]string{
+				"username": "alice", "product": "apple", "review": "bad",
+			})
+
+			// Inserting values into the new schema that violate uniqueness should fail.
+			MustInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "bob", "product": "orange", "review": "good",
+			})
+			MustNotInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "bob", "product": "orange", "review": "bad",
+			})
+		},
+		afterRollback: func(t *testing.T, db *sql.DB) {
+			// The unique index has been dropped from the the underlying table.
+			idxName := migrations.IndexName("reviews", []string{"username", "product"})
+			IndexMustNotExist(t, db, "public", "reviews", idxName)
+		},
+		afterComplete: func(t *testing.T, db *sql.DB) {
+			// The unique constraint has been created on the underlying table.
+			constraintName := migrations.UniqueConstraintName("reviews", []string{"username", "product"})
+			ConstraintMustExist(t, db, "public", "reviews", constraintName)
+
+			// Inserting values into the new schema that violate uniqueness should fail.
+			MustInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "carl", "product": "banana", "review": "good",
+			})
+			MustNotInsert(t, db, "public", "02_set_unique", "reviews", map[string]string{
+				"username": "carl", "product": "banana", "review": "bad",
+			})
+		},
+	}})
+}