From 67f935d1305778b010a1fd4b37854355c5cb18e7 Mon Sep 17 00:00:00 2001
From: everpcpc <everpcpc@users.noreply.github.com>
Date: Wed, 28 Aug 2024 14:43:25 +0800
Subject: [PATCH] fix: test format_null_as_str with conn (#139)

* fix: test format_null_as_str with conn

* z
---
 tests/docker-compose.yml |  2 +-
 tests/nullable_test.go   | 23 +++++++++++++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml
index d3ede21..1493207 100644
--- a/tests/docker-compose.yml
+++ b/tests/docker-compose.yml
@@ -6,7 +6,7 @@ services:
     volumes:
       - ./data:/data
   databend:
-    image: datafuselabs/databend:nightly
+    image: datafuselabs/databend
     environment:
       - QUERY_DEFAULT_USER=databend
       - QUERY_DEFAULT_PASSWORD=databend
diff --git a/tests/nullable_test.go b/tests/nullable_test.go
index 0c0b1cc..9aaf283 100644
--- a/tests/nullable_test.go
+++ b/tests/nullable_test.go
@@ -1,6 +1,7 @@
 package tests
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 )
@@ -16,18 +17,19 @@ func (s *DatabendTestSuite) TestNullable() {
 	s.r.Equal([][]interface{}{{int64(1), nil, nil, "NULL", "NULL", nil, nil, nil, nil}}, result)
 	s.r.NoError(rows.Close())
 
-	_, err = s.db.Exec("SET GLOBAL format_null_as_str=0")
+	ctx := context.TODO()
+	conn, err := s.db.Conn(ctx)
 	s.r.Nil(err)
 
-	rows, err = s.db.Query(fmt.Sprintf("SELECT * FROM %s", s.table))
+	_, err = conn.ExecContext(ctx, "SET format_null_as_str=0")
+	s.r.Nil(err)
+
+	rows, err = conn.QueryContext(ctx, fmt.Sprintf("SELECT * FROM %s", s.table))
 	s.r.Nil(err)
 	result, err = scanValues(rows)
 	s.r.Nil(err)
 	s.r.Equal([][]interface{}{{int64(1), nil, nil, nil, nil, nil, nil, nil, nil}}, result)
 	s.r.NoError(rows.Close())
-
-	_, err = s.db.Exec("UNSET format_null_as_str")
-	s.r.Nil(err)
 }
 
 func (s *DatabendTestSuite) TestQueryNullAsStr() {
@@ -40,16 +42,17 @@ func (s *DatabendTestSuite) TestQueryNullAsStr() {
 }
 
 func (s *DatabendTestSuite) TestQueryNull() {
-	_, err := s.db.Exec("SET GLOBAL format_null_as_str=0")
+	ctx := context.TODO()
+	conn, err := s.db.Conn(ctx)
 	s.r.Nil(err)
 
-	row := s.db.QueryRow("SELECT NULL")
+	_, err = conn.ExecContext(ctx, "SET format_null_as_str=0")
+	s.r.Nil(err)
+
+	row := conn.QueryRowContext(ctx, "SELECT NULL")
 	var val sql.NullString
 	err = row.Scan(&val)
 	s.r.Nil(err)
 	s.r.False(val.Valid)
 	s.r.Equal("", val.String)
-
-	_, err = s.db.Exec("UNSET format_null_as_str")
-	s.r.Nil(err)
 }