Golang server.StartTestServer函數代碼示例

func newKVNative(b *testing.B) kvInterface {
	enableTracing := tracing.Disable()
	s := server.StartTestServer(b)

	// TestServer.DB() returns the TxnCoordSender wrapped client. But that isn't
	// a fair comparison with SQL as we want these client requests to be sent
	// over the network.
	sender, err := client.NewSender(
		rpc.NewContext(&base.Context{
			User:       security.NodeUser,
			SSLCA:      filepath.Join(security.EmbeddedCertsDir, security.EmbeddedCACert),
			SSLCert:    filepath.Join(security.EmbeddedCertsDir, "node.crt"),
			SSLCertKey: filepath.Join(security.EmbeddedCertsDir, "node.key"),
		}, nil, s.Stopper()),
		s.ServingAddr())
	if err != nil {
		b.Fatal(err)
	}

	return &kvNative{
		db: client.NewDB(sender),
		doneFn: func() {
			s.Stop()
			enableTracing()
		},
	}
}

func TestPGWireDBName(t *testing.T) {
	defer leaktest.AfterTest(t)()

	s := server.StartTestServer(t)
	defer s.Stop()

	pgURL, cleanupFn := sqlutils.PGUrl(t, s.ServingAddr(), security.RootUser, "TestPGWireDBName")
	pgURL.Path = "foo"
	defer cleanupFn()
	{
		db, err := gosql.Open("postgres", pgURL.String())
		if err != nil {
			t.Fatal(err)
		}
		defer db.Close()

		if _, err := db.Exec(`CREATE DATABASE foo`); err != nil {
			t.Fatal(err)
		}

		if _, err := db.Exec(`CREATE TABLE bar (i INT PRIMARY KEY)`); err != nil {
			t.Fatal(err)
		}
	}
	db, err := gosql.Open("postgres", pgURL.String())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	if _, err := db.Exec(`INSERT INTO bar VALUES ($1)`, 1); err != nil {
		t.Fatal(err)
	}
}

// TestRangeLookupWithOpenTransaction verifies that range lookups are
// done in such a way (e.g. using inconsistent reads) that they
// proceed in the event that a write intent is extant at the meta
// index record being read.
func TestRangeLookupWithOpenTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()
	db := createTestClient(t, s.Stopper(), s.ServingAddr())

	// Create an intent on the meta1 record by writing directly to the
	// engine.
	key := testutils.MakeKey(keys.Meta1Prefix, roachpb.KeyMax)
	now := s.Clock().Now()
	txn := roachpb.NewTransaction("txn", roachpb.Key("foobar"), 0, roachpb.SERIALIZABLE, now, 0)
	if err := engine.MVCCPutProto(s.Ctx.Engines[0], nil, key, now, txn, &roachpb.RangeDescriptor{}); err != nil {
		t.Fatal(err)
	}

	// Now, with an intent pending, attempt (asynchronously) to read
	// from an arbitrary key. This will cause the distributed sender to
	// do a range lookup, which will encounter the intent. We're
	// verifying here that the range lookup doesn't fail with a write
	// intent error. If it did, it would go into a deadloop attempting
	// to push the transaction, which in turn requires another range
	// lookup, etc, ad nauseam.
	if _, err := db.Get("a"); err != nil {
		t.Fatal(err)
	}
}

// TestKVDBInternalMethods verifies no internal methods are available
// HTTP DB interface.
func TestKVDBInternalMethods(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()

	testCases := []struct {
		args  proto.Request
		reply proto.Response
	}{
		{&proto.InternalRangeLookupRequest{}, &proto.InternalRangeLookupResponse{}},
		{&proto.InternalGCRequest{}, &proto.InternalGCResponse{}},
		{&proto.InternalHeartbeatTxnRequest{}, &proto.InternalHeartbeatTxnResponse{}},
		{&proto.InternalPushTxnRequest{}, &proto.InternalPushTxnResponse{}},
		{&proto.InternalResolveIntentRequest{}, &proto.InternalResolveIntentResponse{}},
		{&proto.InternalResolveIntentRangeRequest{}, &proto.InternalResolveIntentRangeResponse{}},
		{&proto.InternalMergeRequest{}, &proto.InternalMergeResponse{}},
		{&proto.InternalTruncateLogRequest{}, &proto.InternalTruncateLogResponse{}},
	}
	// Verify non-public methods experience bad request errors.
	db := createTestClient(t, s.ServingAddr())
	for i, test := range testCases {
		test.args.Header().Key = proto.Key("a")
		if proto.IsRange(test.args) {
			test.args.Header().EndKey = test.args.Header().Key.Next()
		}
		b := &client.Batch{}
		b.InternalAddCall(proto.Call{Args: test.args, Reply: test.reply})
		err := db.Run(b)
		if err == nil {
			t.Errorf("%d: unexpected success calling %s", i, test.args.Method())
		} else if !strings.Contains(err.Error(), "404 Not Found") {
			t.Errorf("%d: expected 404; got %s", i, err)
		}
	}
}

func TestBadRequest(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	db := createTestClient(t, s.Stopper(), s.ServingAddr())
	defer s.Stop()

	// Write key "a".
	if pErr := db.Put("a", "value"); pErr != nil {
		t.Fatal(pErr)
	}

	if _, pErr := db.Scan("a", "a", 0); !testutils.IsPError(pErr, "truncation resulted in empty batch") {
		t.Fatalf("unexpected error on scan with startkey == endkey: %v", pErr)
	}

	if _, pErr := db.ReverseScan("a", "a", 0); !testutils.IsPError(pErr, "truncation resulted in empty batch") {
		t.Fatalf("unexpected pError on reverse scan with startkey == endkey: %v", pErr)
	}

	if pErr := db.DelRange("x", "a"); !testutils.IsPError(pErr, "truncation resulted in empty batch") {
		t.Fatalf("unexpected error on deletion on [x, a): %v", pErr)
	}

	if pErr := db.DelRange("", "z"); !testutils.IsPError(pErr, "must be greater than LocalMax") {
		t.Fatalf("unexpected error on deletion on [KeyMin, z): %v", pErr)
	}
}

// TestSingleRangeReverseScan verifies that ReverseScan gets the right results
// on a single range.
func TestSingleRangeReverseScan(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()
	db := initReverseScanTestEnv(s, t)

	// Case 1: Request.EndKey is in the middle of the range.
	if rows, pErr := db.ReverseScan("b", "d", 0); pErr != nil {
		t.Fatalf("unexpected pError on ReverseScan: %s", pErr)
	} else if l := len(rows); l != 2 {
		t.Errorf("expected 2 rows; got %d", l)
	}
	// Case 2: Request.EndKey is equal to the EndKey of the range.
	if rows, pErr := db.ReverseScan("e", "g", 0); pErr != nil {
		t.Fatalf("unexpected error on ReverseScan: %s", pErr)
	} else if l := len(rows); l != 2 {
		t.Errorf("expected 2 rows; got %d", l)
	}
	// Case 3: Test roachpb.TableDataMin. Expected to return "g" and "h".
	wanted := 2
	if rows, pErr := db.ReverseScan("g", keys.TableDataMin, 0); pErr != nil {
		t.Fatalf("unexpected error on ReverseScan: %s", pErr)
	} else if l := len(rows); l != wanted {
		t.Errorf("expected %d rows; got %d", wanted, l)
	}
	// Case 4: Test keys.SystemMax
	// This span covers the system DB keys. Note sql.GetInitialSystemValues
	// returns one key before keys.SystemMax, but our scan is including one key
	// (\xffa) created for the test.
	if rows, pErr := db.ReverseScan(keys.SystemMax, "b", 0); pErr != nil {
		t.Fatalf("unexpected error on ReverseScan: %s", pErr)
	} else if l := len(rows); l != 1 {
		t.Errorf("expected 1 row; got %d", l)
	}
}

// TestClientGetAndPutProto verifies gets and puts of protobufs using the
// client's convenience methods.
func TestClientGetAndPutProto(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()
	db := createTestClient(t, s.Stopper(), s.ServingAddr())

	zoneConfig := &config.ZoneConfig{
		ReplicaAttrs: []roachpb.Attributes{
			{Attrs: []string{"dc1", "mem"}},
			{Attrs: []string{"dc2", "mem"}},
		},
		RangeMinBytes: 1 << 10, // 1k
		RangeMaxBytes: 1 << 18, // 256k
	}

	key := roachpb.Key(testUser + "/zone-config")
	if pErr := db.Put(key, zoneConfig); pErr != nil {
		t.Fatalf("unable to put proto: %s", pErr)
	}

	readZoneConfig := &config.ZoneConfig{}
	if pErr := db.GetProto(key, readZoneConfig); pErr != nil {
		t.Fatalf("unable to get proto: %v", pErr)
	}
	if !proto.Equal(zoneConfig, readZoneConfig) {
		t.Errorf("expected %+v, but found %+v", zoneConfig, readZoneConfig)
	}
}

func TestMultiRangeBatchBoundedScans(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s := server.StartTestServer(t)
	defer s.Stop()
	db := setupMultipleRanges(t, s, "a", "b", "c", "d", "e", "f")
	for _, key := range []string{"a", "aa", "aaa", "b", "bb", "cc", "d", "dd", "ff"} {
		if err := db.Put(key, "value"); err != nil {
			t.Fatal(err)
		}
	}

	b := db.NewBatch()
	b.Scan("aaa", "dd", 3)
	b.Scan("a", "z", 2)
	b.Scan("cc", "ff", 3)
	if err := db.Run(b); err != nil {
		t.Fatal(err)
	}

	checkScanResults(t, b.Results, [][]string{
		{"aaa", "b", "bb"},
		{"a", "aa"},
		{"cc", "d", "dd"},
	})
}

func TestPGWireMetrics(t *testing.T) {
	defer leaktest.AfterTest(t)

	s := server.StartTestServer(t)
	defer s.Stop()

	// Setup pgwire client.
	pgUrl, cleanupFn := sqlutils.PGUrl(t, s, security.RootUser, "TestPGWireMetrics")
	defer cleanupFn()

	const minbytes = 20

	// Make sure we're starting at 0.
	if _, _, err := checkPGWireMetrics(s, 0, 0, 0, 0); err != nil {
		t.Fatal(err)
	}

	// A single query should give us some I/O.
	if err := trivialQuery(pgUrl); err != nil {
		t.Fatal(err)
	}
	bytesIn, bytesOut, err := checkPGWireMetrics(s, minbytes, minbytes, 300, 300)
	if err != nil {
		t.Fatal(err)
	}
	if err := trivialQuery(pgUrl); err != nil {
		t.Fatal(err)
	}

	// A second query should give us more I/O.
	_, _, err = checkPGWireMetrics(s, bytesIn+minbytes, bytesOut+minbytes, 300, 300)
	if err != nil {
		t.Fatal(err)
	}
}

// TestKVDBEndTransactionWithTriggers verifies that triggers are
// disallowed on call to EndTransaction.
func TestKVDBEndTransactionWithTriggers(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()

	db := createTestClient(t, s.ServingAddr())
	err := db.Txn(func(txn *client.Txn) error {
		// Make an EndTransaction request which would fail if not
		// stripped. In this case, we set the start key to "bar" for a
		// split of the default range; start key must be "" in this case.
		b := &client.Batch{}
		b.Put("foo", "only here to make this a rw transaction")
		b.InternalAddCall(proto.Call{
			Args: &proto.EndTransactionRequest{
				RequestHeader: proto.RequestHeader{Key: proto.Key("foo")},
				Commit:        true,
				InternalCommitTrigger: &proto.InternalCommitTrigger{
					SplitTrigger: &proto.SplitTrigger{
						UpdatedDesc: proto.RangeDescriptor{StartKey: proto.Key("bar")},
					},
				},
			},
			Reply: &proto.EndTransactionResponse{},
		})
		return txn.Run(b)
	})
	if err == nil {
		t.Errorf("expected 400 bad request error on commit")
	}
}

// TestHTTPAuthentication tests authentication for the KV http endpoint.
func TestHTTPAuthentication(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()

	// createTestClient creates a "root" client.
	db := createTestClient(t, s.ServingAddr())

	// We call Run() on the client which lets us build our own request,
	// specifying the user.
	arg := &proto.PutRequest{}
	arg.Header().Key = proto.Key("a")
	arg.Header().User = security.RootUser
	reply := &proto.PutResponse{}
	b := &client.Batch{}
	b.InternalAddCall(proto.Call{Args: arg, Reply: reply})
	err := db.Run(b)
	if err != nil {
		t.Fatal(err)
	}

	// Try again, but this time with arg.User = "foo".
	arg.Header().User = "foo"
	b = &client.Batch{}
	b.InternalAddCall(proto.Call{Args: arg, Reply: reply})
	err = db.Run(b)
	if err == nil {
		t.Fatal("Expected error!")
	}
}

// TestAuthentication tests authentication for the KV endpoint.
func TestAuthentication(t *testing.T) {
	defer leaktest.AfterTest(t)()
	s := server.StartTestServer(t)
	defer s.Stop()

	var b1 client.Batch
	b1.Put("a", "b")

	// Create a node user client and call Run() on it which lets us build our own
	// request, specifying the user.
	db1 := createTestClientForUser(t, s.Stopper(), s.ServingAddr(), security.NodeUser)
	if err := db1.Run(&b1); err != nil {
		t.Fatal(err)
	}

	var b2 client.Batch
	b2.Put("c", "d")

	// Try again, but this time with certs for a non-node user (even the root
	// user has no KV permissions).
	db2 := createTestClientForUser(t, s.Stopper(), s.ServingAddr(), security.RootUser)
	if err := db2.Run(&b2); !testutils.IsError(err, "is not allowed") {
		t.Fatal(err)
	}
}

// TestNoSequenceCachePutOnRangeMismatchError verifies that the
// sequence cache is not updated with RangeKeyMismatchError. This is a
// higher-level version of TestSequenceCacheShouldCache.
func TestNoSequenceCachePutOnRangeMismatchError(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()
	db := setupMultipleRanges(t, s, "b", "c")

	// The requests in the transaction below will be chunked and
	// sent to replicas in the following way:
	// 1) A batch request containing a BeginTransaction and a
	//    put on "a" are sent to a replica owning range ["a","b").
	// 2) A next batch request containing a put on "b" and a put
	//    on "c" are sent to a replica owning range ["b","c").
	//   (The range cache has a stale range descriptor.)
	// 3) The put request on "c" causes a RangeKeyMismatchError.
	// 4) The dist sender re-sends a request to the same replica.
	//    This time the request contains only the put on "b" to the
	//    same replica.
	// 5) The command succeeds since the sequence cache has not yet been updated.
	epoch := 0
	if pErr := db.Txn(func(txn *client.Txn) *roachpb.Error {
		epoch++
		b := txn.NewBatch()
		b.Put("a", "val")
		b.Put("b", "val")
		b.Put("c", "val")
		return txn.CommitInBatch(b)
	}); pErr != nil {
		t.Errorf("unexpected error on transactional Puts: %s", pErr)
	}

	if epoch != 1 {
		t.Errorf("unexpected epoch; the txn must not be retried, but got %d retries", epoch)
	}
}

// TestReverseScanWithSplitAndMerge verifies that ReverseScan gets the right results
// across multiple ranges while range splits and merges happen.
func TestReverseScanWithSplitAndMerge(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()
	db := initReverseScanTestEnv(s, t)

	// Case 1: An encounter with a range split.
	// Split the range ["b", "e") at "c".
	if pErr := db.AdminSplit("c"); pErr != nil {
		t.Fatal(pErr)
	}

	// The ReverseScan will run into a stale descriptor.
	if rows, pErr := db.ReverseScan("a", "d", 0); pErr != nil {
		t.Fatalf("unexpected error on ReverseScan: %s", pErr)
	} else if l := len(rows); l != 3 {
		t.Errorf("expected 3 rows; got %d", l)
	}

	// Case 2: encounter with range merge .
	// Merge the range ["e", "g") and ["g", "\xff\xff") .
	if pErr := db.AdminMerge("e"); pErr != nil {
		t.Fatal(pErr)
	}
	if rows, pErr := db.ReverseScan("d", "g", 0); pErr != nil {
		t.Fatalf("unexpected error on ReverseScan: %s", pErr)
	} else if l := len(rows); l != 3 {
		t.Errorf("expected 3 rows; got %d", l)
	}
}

func TestSend(t *testing.T) {
	defer leaktest.AfterTest(t)
	s := server.StartTestServer(t)
	defer s.Stop()
	sender, err := newHTTPSender(s.ServingAddr(), testutils.NewRootTestBaseContext(), defaultRetryOptions)
	if err != nil {
		log.Fatalf("Couldn't create HTTPSender for server:(%s)", s.ServingAddr())
	}
	testCases := []struct {
		req   string
		reply string
	}{
		{"ping", "ping"},
		{"default", "default"},
	}
	for _, test := range testCases {
		request := &sqlwire.Request{Sql: test.req}
		call := sqlwire.Call{Args: request, Reply: &sqlwire.Response{}}
		sender.Send(context.TODO(), call)
		reply := *call.Reply.Results[0].Rows[0].Values[0].StringVal
		if reply != test.reply {
			log.Fatalf("Server sent back reply: %s", reply)
		}
	}
}