From 21ffb64467730dc7c59ea20e7556d92d46082544 Mon Sep 17 00:00:00 2001
From: James Cor <jcor@ucsd.edu>
Date: Mon, 8 Sep 2025 15:52:28 -0700
Subject: [PATCH] fix dolt test run in server mode (#9811)

Co-authored-by: James Cor <james@dolthub.com>
---
 .../sqle/dtablefunctions/dolt_test_run.go     | 22 +++++++++++++-----
 integration-tests/bats/dolt-test-run.bats     | 23 +++++++++++++++++++
 2 files changed, 39 insertions(+), 6 deletions(-)
 create mode 100644 integration-tests/bats/dolt-test-run.bats

diff --git a/go/libraries/doltcore/sqle/dtablefunctions/dolt_test_run.go b/go/libraries/doltcore/sqle/dtablefunctions/dolt_test_run.go
index eb42788103..343be685e1 100644
--- a/go/libraries/doltcore/sqle/dtablefunctions/dolt_test_run.go
+++ b/go/libraries/doltcore/sqle/dtablefunctions/dolt_test_run.go
@@ -16,6 +16,7 @@ package dtablefunctions
 
 import (
 	"fmt"
+	"io"
 	"strings"
 
 	gms "github.com/dolthub/go-mysql-server"
@@ -191,7 +192,7 @@ func (trtf *TestsRunTableFunction) RowIter(_ *sql.Context, _ sql.Row) (sql.RowIt
 			return nil, err
 		}
 
-		for _, row := range *testRows {
+		for _, row := range testRows {
 			result, err := trtf.queryAndAssert(row)
 			if err != nil {
 				return nil, err
@@ -254,7 +255,7 @@ func (trtf *TestsRunTableFunction) queryAndAssert(row sql.Row) (result testResul
 	return result, nil
 }
 
-func (trtf *TestsRunTableFunction) getDoltTestsData(arg string) (*[]sql.Row, error) {
+func (trtf *TestsRunTableFunction) getDoltTestsData(arg string) ([]sql.Row, error) {
 	var queries []string
 
 	if arg == "*" {
@@ -280,12 +281,21 @@ func (trtf *TestsRunTableFunction) getDoltTestsData(arg string) (*[]sql.Row, err
 		if err != nil {
 			return nil, err
 		}
-		rows, err := sql.RowIterToRows(trtf.ctx, iter)
-		if err != nil {
-			return nil, err
+		// Calling iter.Close(ctx) will cause TrackedRowIter to cancel the context, causing problems when running with
+		// dolt sql-server. Since we only support `SELECT...` queries anyway, it's not necessary to Close() the iter.
+		var rows []sql.Row
+		for {
+			row, rErr := iter.Next(trtf.ctx)
+			if rErr == io.EOF {
+				break
+			}
+			if rErr != nil {
+				return nil, rErr
+			}
+			rows = append(rows, row)
 		}
 		if len(rows) > 0 {
-			return &rows, nil
+			return rows, nil
 		}
 	}
 	return nil, fmt.Errorf("could not find tests for argument: %s", arg)
diff --git a/integration-tests/bats/dolt-test-run.bats b/integration-tests/bats/dolt-test-run.bats
new file mode 100644
index 0000000000..d4877a43a3
--- /dev/null
+++ b/integration-tests/bats/dolt-test-run.bats
@@ -0,0 +1,23 @@
+#!/usr/bin/env bats
+load $BATS_TEST_DIRNAME/helper/common.bash
+
+setup() {
+    if [ "$SQL_ENGINE" = "remote-engine" ]; then
+      skip "This test tests remote connections directly, SQL_ENGINE is not needed."
+    fi
+    setup_common
+}
+
+teardown() {
+    stop_sql_server 1 && sleep 0.5
+    teardown_common
+}
+
+@test "dolt-test-run: sanity test on sql-server" {
+    start_sql_server
+
+    dolt sql -q "insert into dolt_tests values ('test', 'test', 'select 1', 'expected_rows', '==', '1');"
+    run dolt sql -q "select * from dolt_test_run()"
+    [ $status -eq 0 ]
+    [[ $output =~ "| test      | test            | select 1 | PASS   |         |" ]] || false
+}