diff --git a/server/conn.go b/server/conn.go index 0277bd40a49be..f74c00a5550cc 100644 --- a/server/conn.go +++ b/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() } diff --git a/server/mock_conn.go b/server/mock_conn.go new file mode 100644 index 0000000000000..2d3c8d581466d --- /dev/null +++ b/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, + } +} diff --git a/server/mock_conn_test.go b/server/mock_conn_test.go new file mode 100644 index 0000000000000..9eae5ad1003da --- /dev/null +++ b/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) +}