tidb mock_conn 源码

tidb mock_conn 代码

文件路径：/server/mock_conn.go

// 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"
	"flag"
	"math/rand"
	"testing"

	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/parser/auth"
	tmysql "github.com/pingcap/tidb/parser/mysql"
	"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()
	// ID returns the connection ID.
	ID() uint64
}

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.getCtx()
}

// 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())
}

// ID implements MockConn.ID
func (mc *mockConn) ID() uint64 {
	return mc.clientConn.connectionID
}

// CreateMockServer creates a mock server.
func CreateMockServer(t *testing.T, store kv.Storage) *Server {
	if !RunInGoTest {
		// If CreateMockServer is called in another package, RunInGoTest is not initialized.
		RunInGoTest = flag.Lookup("test.v") != nil || flag.Lookup("check.v") != nil
	}
	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, server *Server) MockConn {
	connID := rand.Uint64()
	tc, err := server.driver.OpenCtx(connID, 0, uint8(tmysql.DefaultCollationID), "", nil)
	require.NoError(t, err)

	cc := &clientConn{
		connectionID: connID,
		server:       server,
		salt:         []byte{},
		collation:    tmysql.DefaultCollationID,
		alloc:        arena.NewAllocator(1024),
		chunkAlloc:   chunk.NewAllocator(),
		pkt: &packetIO{
			bufWriter: bufio.NewWriter(bytes.NewBuffer(nil)),
		},
	}
	cc.setCtx(tc)
	cc.server.rwlock.Lock()
	server.clients[cc.connectionID] = cc
	cc.server.rwlock.Unlock()
	tc.Session.SetSessionManager(server)
	err = tc.Session.Auth(&auth.UserIdentity{Username: "root", Hostname: "localhost"}, nil, nil)
	require.NoError(t, err)
	return &mockConn{
		clientConn: cc,
		t:          t,
	}
}

相关信息

tidb 源码目录

相关文章

tidb buffered_read_conn 源码

tidb column 源码

tidb conn 源码

tidb conn_stmt 源码

tidb driver 源码

tidb driver_tidb 源码

tidb http_handler 源码

tidb http_status 源码

tidb optimize_trace 源码

tidb packetio 源码