Fix tests to make them actually drain iterators

go/libraries/doltcore/sqle/alterschema_test.go

@@ -16,7 +16,7 @@ package sqle
 
 import (
 	"context"
-	"errors"
+	goerrors "errors"
 	"fmt"
 	"testing"
 
@@ -25,6 +25,7 @@ import (
 	"github.com/dolthub/vitess/go/sqltypes"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gopkg.in/src-d/go-errors.v1"
 
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/dtestutils"
@@ -445,10 +446,10 @@ func TestDropColumnUsedByIndex(t *testing.T) {
 
 func TestDropPks(t *testing.T) {
 	var dropTests = []struct {
-		name      string
-		setup     []string
-		exit      int
-		fkIdxName string
+		name           string
+		setup          []string
+		expectedErr    *errors.Kind
+		fkIdxName      string
 	}{
 		{
 			name: "no error on drop pk",
@@ -456,7 +457,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id))",
 				"insert into parent values (1,1,1),(2,2,2)",
 			},
-			exit: 0,
 		},
 		{
 			name: "no error if backup index",
@@ -464,7 +464,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id), key `backup` (id))",
 				"create table child (id int, name varchar(1), age int, primary key (id), constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -473,7 +472,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age), key `backup` (age))",
 				"create table child (id int, name varchar(1), age int, primary key (id), constraint `fk` foreign key (age) references parent (age))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -482,7 +480,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age), key `backup` (id, age))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id, age) references parent (id, age))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -491,7 +488,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age, name), key `backup` (id, age))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id, age) references parent (id, age))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -500,7 +496,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age), key `backup` (id))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -509,7 +504,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age), key `bad_backup1` (age, id), key `bad_backup2` (age), key `backup` (id, age, name))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -518,7 +512,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age), key `bad_backup` (age, id), key `backup1` (id), key `backup2` (id, age, name))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      0,
 			fkIdxName: "backup1",
 		},
 		{
@@ -527,7 +520,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id, age), key `bad_backup` (age, id), key `backup1` (id, age, name), unique key `backup2` (id, age))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      0,
 			fkIdxName: "backup2",
 		},
 		{
@@ -536,7 +528,6 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, other int, primary key (id, age, name), key `backup` (id, age, other))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id, age) references parent (id, age))",
 			},
-			exit:      0,
 			fkIdxName: "backup",
 		},
 		{
@@ -545,7 +536,7 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id))",
 				"create table child (id int, name varchar(1), age int, primary key (id), constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      1,
+			expectedErr:  sql.ErrCantDropIndex,
 			fkIdxName: "id",
 		},
 		{
@@ -554,7 +545,7 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, primary key (id), key `bad_backup2` (age))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      1,
+			expectedErr:  sql.ErrCantDropIndex,
 			fkIdxName: "id",
 		},
 		{
@@ -563,7 +554,7 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, constraint `primary` primary key (id), key `backup` (age, id))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (id) references parent (id))",
 			},
-			exit:      1,
+			expectedErr:  sql.ErrCantDropIndex,
 			fkIdxName: "id",
 		},
 		{
@@ -572,7 +563,7 @@ func TestDropPks(t *testing.T) {
 				"create table parent (id int, name varchar(1), age int, constraint `primary` primary key (age, id), key `backup` (age, name))",
 				"create table child (id int, name varchar(1), age int, constraint `fk` foreign key (age, id) references parent (age,  id))",
 			},
-			exit:      1,
+			expectedErr:  sql.ErrCantDropIndex,
 			fkIdxName: "ageid",
 		},
 	}
@@ -598,11 +589,14 @@ func TestDropPks(t *testing.T) {
 			}
 
 			drop := "alter table parent drop primary key"
-			_, _, err = engine.Query(sqlCtx, drop)
-			switch tt.exit {
-			case 1:
+			_, iter, err := engine.Query(sqlCtx, drop)
+			require.NoError(t, err)
+
+			err = drainIter(sqlCtx, iter)
+			if tt.expectedErr != nil {
 				require.Error(t, err)
-			default:
+				assert.True(t, tt.expectedErr.Is(err), "Expected error of type %s but got %s", tt.expectedErr, err)
+			} else {
 				require.NoError(t, err)
 			}
 
@@ -771,7 +765,7 @@ func TestNewPkOrdinals(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			res, err := modifyPkOrdinals(oldSch, tt.newSch)
 			if tt.err != nil {
-				require.True(t, errors.Is(err, tt.err))
+				require.True(t, goerrors.Is(err,  tt.err))
 			} else {
 				require.Equal(t, res, tt.expPkOrdinals)
 			}