Extending script queries to allow easier testing of warning messages

enginetest/enginetests.go

@@ -1208,7 +1208,11 @@ func TestScriptWithEngine(t *testing.T, e *sqle.Engine, harness Harness, script
 				AssertErr(t, e, harness, assertion.Query, nil, assertion.ExpectedErrStr)
 			})
 		} else if assertion.ExpectedWarning != 0 {
-			AssertWarningAndTestQuery(t, e, nil, harness, assertion.Query, assertion.Expected, nil, assertion.ExpectedWarning)
+			t.Run(assertion.Query, func(t *testing.T) {
+				AssertWarningAndTestQuery(t, e, nil, harness, assertion.Query,
+					assertion.Expected, nil, assertion.ExpectedWarning, assertion.ExpectedWarningsCount,
+					assertion.ExpectedWarningMessageSubstring, assertion.SkipResultsCheck)
+			})
 		} else {
 			TestQuery(t, harness, e, assertion.Query, assertion.Expected, nil, nil)
 		}
@@ -1256,7 +1260,9 @@ func TestTransactionScriptWithEngine(t *testing.T, e *sqle.Engine, harness Harne
 			} else if assertion.ExpectedErrStr != "" {
 				AssertErrWithCtx(t, e, clientSession, assertion.Query, nil, assertion.ExpectedErrStr)
 			} else if assertion.ExpectedWarning != 0 {
-				AssertWarningAndTestQuery(t, e, nil, harness, assertion.Query, assertion.Expected, nil, assertion.ExpectedWarning)
+				AssertWarningAndTestQuery(t, e, nil, harness, assertion.Query, assertion.Expected,
+					nil, assertion.ExpectedWarning, assertion.ExpectedWarningsCount,
+					assertion.ExpectedWarningMessageSubstring, false)
 			} else {
 				TestQueryWithContext(t, clientSession, e, assertion.Query, assertion.Expected, nil, nil)
 			}
@@ -2235,7 +2241,9 @@ func TestCreateDatabase(t *testing.T, harness Harness) {
 	t.Run("CREATE DATABASE error handling", func(t *testing.T) {
 		AssertErr(t, e, harness, "CREATE DATABASE mydb", sql.ErrDatabaseExists)
 
-		AssertWarningAndTestQuery(t, e, nil, harness, "CREATE DATABASE IF NOT EXISTS mydb", []sql.Row{{sql.OkResult{RowsAffected: 1}}}, nil, mysql.ERDbCreateExists)
+		AssertWarningAndTestQuery(t, e, nil, harness, "CREATE DATABASE IF NOT EXISTS mydb",
+			[]sql.Row{{sql.OkResult{RowsAffected: 1}}}, nil, mysql.ERDbCreateExists,
+			-1, "", false)
 	})
 }
 
@@ -2401,7 +2409,9 @@ func TestDropDatabase(t *testing.T, harness Harness) {
 		// The test setup sets a database name, which interferes with DROP DATABASE tests
 		ctx := NewContext(harness)
 		TestQueryWithContext(t, ctx, e, "DROP DATABASE mydb", []sql.Row{{sql.OkResult{RowsAffected: 1}}}, nil, nil)
-		AssertWarningAndTestQuery(t, e, ctx, harness, "DROP DATABASE IF EXISTS mydb", []sql.Row{{sql.OkResult{RowsAffected: 0}}}, nil, mysql.ERDbDropExists)
+		AssertWarningAndTestQuery(t, e, ctx, harness, "DROP DATABASE IF EXISTS mydb",
+			[]sql.Row{{sql.OkResult{RowsAffected: 0}}}, nil, mysql.ERDbDropExists,
+			-1, "", false)
 
 		TestQueryWithContext(t, ctx, e, "CREATE DATABASE testdb", []sql.Row{{sql.OkResult{RowsAffected: 1}}}, nil, nil)
 
@@ -2417,7 +2427,9 @@ func TestDropDatabase(t *testing.T, harness Harness) {
 		require.Error(t, err)
 		require.True(t, sql.ErrDatabaseNotFound.Is(err), "Expected error of type %s but got %s", sql.ErrDatabaseNotFound, err)
 
-		AssertWarningAndTestQuery(t, e, ctx, harness, "DROP DATABASE IF EXISTS testdb", []sql.Row{{sql.OkResult{RowsAffected: 0}}}, nil, mysql.ERDbDropExists)
+		AssertWarningAndTestQuery(t, e, ctx, harness, "DROP DATABASE IF EXISTS testdb",
+			[]sql.Row{{sql.OkResult{RowsAffected: 0}}}, nil, mysql.ERDbDropExists,
+			-1, "", false)
 	})
 }
 
@@ -4164,28 +4176,41 @@ func AssertWarningAndTestQuery(
 	expected []sql.Row,
 	expectedCols []*sql.Column,
 	expectedCode int,
+	expectedWarningsCount int,
+	expectedWarningMessageSubstring string,
+	skipResultsCheck bool,
 ) {
 	require := require.New(t)
 	if ctx == nil {
 		ctx = NewContext(harness)
 	}
+	ctx.ClearWarnings()
+
 	sch, iter, err := e.Query(ctx, query)
 	require.NoError(err, "Unexpected error for query %s", query)
 
 	rows, err := sql.RowIterToRows(ctx, sch, iter)
 	require.NoError(err, "Unexpected error for query %s", query)
 
-	condition := false
-	for _, warning := range ctx.Warnings() {
-		if warning.Code == expectedCode {
-			condition = true
-			break
+	if expectedWarningsCount > 0 {
+		assert.Equal(t, expectedWarningsCount, len(ctx.Warnings()))
+	}
+
+	if expectedCode > 0 {
+		for _, warning := range ctx.Warnings() {
+			assert.Equal(t, expectedCode, warning.Code, "Unexpected warning code")
 		}
 	}
 
-	assert.True(t, condition)
+	if len(expectedWarningMessageSubstring) > 0 {
+		for _, warning := range ctx.Warnings() {
+			assert.Contains(t, warning.Message, expectedWarningMessageSubstring, "Unexpected warning message")
+		}
+	}
 
-	checkResults(t, require, expected, expectedCols, sch, rows, query)
+	if !skipResultsCheck {
+		checkResults(t, require, expected, expectedCols, sch, rows, query)
+	}
 }
 
 type customFunc struct {

enginetest/script_queries.go

@@ -47,8 +47,20 @@ type ScriptTestAssertion struct {
 	// such as the use of the SIGNAL statement.
 	ExpectedErrStr string
 
-	// ExpectedWarning is used for queries that should generate warnings but not errors.
+	// ExpectedWarning contains the expected warning code when a query generates warnings but not errors.
 	ExpectedWarning int
+
+	// ExpectedWarningsCount is used to test the expected number of warnings generated by a query.
+	// The ExpectedWarning field must be set for warning counts to be checked.
+	ExpectedWarningsCount int
+
+	// ExpectedWarningMessageSubstring is used to test the contents of warning messages generated by a
+	// query. The ExpectedWarning field must be set for warning messages to be checked.
+	ExpectedWarningMessageSubstring string
+
+	// SkipResultsCheck is used to skip assertions on expected Rows returned from a query. This should be used
+	// sparingly, such as in cases where you only want to test warning messages.
+	SkipResultsCheck bool
 }
 
 // ScriptTests are a set of test scripts to run.