server: Add a MockConn to enable testing connections (#30119)

server/conn.go

@@ -310,8 +310,10 @@ func (cc *clientConn) Close() error {
 
 func closeConn(cc *clientConn, connections int) error {
 	metrics.ConnGauge.Set(float64(connections))
-	err := cc.bufReadConn.Close()
-	terror.Log(err)
+	if cc.bufReadConn != nil {
+		err := cc.bufReadConn.Close()
+		terror.Log(err)
+	}
 	if cc.ctx != nil {
 		return cc.ctx.Close()
 	}

server/mock_conn.go

@@ -0,0 +1,106 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package server
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/session"
+	"github.com/pingcap/tidb/util/arena"
+	"github.com/pingcap/tidb/util/chunk"
+	"github.com/stretchr/testify/require"
+)
+
+// MockConn is a mock connection.
+type MockConn interface {
+	// HandleQuery executes a statement
+	HandleQuery(ctx context.Context, sql string) error
+	// Context gets the TiDBContext
+	Context() *TiDBContext
+	// Dispatch executes command according to the command type
+	Dispatch(ctx context.Context, data []byte) error
+	// Close releases resources
+	Close()
+}
+
+type mockConn struct {
+	*clientConn
+	t *testing.T
+}
+
+// HandleQuery implements MockConn.HandleQuery
+func (mc *mockConn) HandleQuery(ctx context.Context, sql string) error {
+	return mc.handleQuery(ctx, sql)
+}
+
+// Context implements MockConn.Context
+func (mc *mockConn) Context() *TiDBContext {
+	return mc.ctx
+}
+
+// Dispatch implements MockConn.Dispatch
+func (mc *mockConn) Dispatch(ctx context.Context, data []byte) error {
+	return mc.dispatch(ctx, data)
+}
+
+// Close implements MockConn.Close
+func (mc *mockConn) Close() {
+	require.NoError(mc.t, mc.clientConn.Close())
+}
+
+// CreateMockServer creates a mock server.
+func CreateMockServer(t *testing.T, store kv.Storage) *Server {
+	tidbdrv := NewTiDBDriver(store)
+	cfg := config.NewConfig()
+	cfg.Socket = ""
+	cfg.Port, cfg.Status.StatusPort = 0, 0
+	cfg.Status.ReportStatus = false
+	cfg.Security.AutoTLS = false
+	server, err := NewServer(cfg, tidbdrv)
+	require.NoError(t, err)
+	return server
+}
+
+// CreateMockConn creates a mock connection together with a session.
+func CreateMockConn(t *testing.T, store kv.Storage, server *Server) MockConn {
+	se, err := session.CreateSession4Test(store)
+	require.NoError(t, err)
+	tc := &TiDBContext{
+		Session: se,
+		stmts:   make(map[int]*TiDBStatement),
+	}
+
+	cc := &clientConn{
+		server:     server,
+		ctx:        tc,
+		salt:       []byte{},
+		collation:  mysql.DefaultCollationID,
+		alloc:      arena.NewAllocator(1024),
+		chunkAlloc: chunk.NewAllocator(),
+		pkt: &packetIO{
+			bufWriter: bufio.NewWriter(bytes.NewBuffer(nil)),
+		},
+	}
+	return &mockConn{
+		clientConn: cc,
+		t:          t,
+	}
+}

server/mock_conn_test.go

@@ -0,0 +1,48 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package server
+
+import (
+	"context"
+	"testing"
+
+	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/testkit"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMockConn(t *testing.T) {
+	t.Parallel()
+
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+	server := CreateMockServer(t, store)
+	defer server.Close()
+	conn := CreateMockConn(t, store, server)
+	defer conn.Close()
+
+	require.NoError(t, conn.HandleQuery(context.Background(), "select 1"))
+	require.Equal(t, "select 1", conn.Context().GetSessionVars().StmtCtx.OriginalSQL)
+
+	require.Error(t, conn.HandleQuery(context.Background(), "select"))
+
+	inBytes := append([]byte{mysql.ComQuery}, []byte("select 1")...)
+	require.NoError(t, conn.Dispatch(context.Background(), inBytes))
+	require.Equal(t, "select 1", conn.Context().GetSessionVars().StmtCtx.OriginalSQL)
+
+	inBytes = append([]byte{mysql.ComStmtPrepare}, []byte("select 1")...)
+	require.NoError(t, conn.Dispatch(context.Background(), inBytes))
+	require.Equal(t, "select 1", conn.Context().GetSessionVars().StmtCtx.OriginalSQL)
+}