From 896180cc26d70887878392457310c0e99e7be7a7 Mon Sep 17 00:00:00 2001 From: Hengfeng Li Date: Tue, 26 Nov 2019 17:04:00 +1100 Subject: [PATCH] spanner: Add resource-based routing This adds a step to retrieve the instance-specific endpoint before creating the session client when creating a new spanner client. It includes the following changes: * Added an extra step to get the instance-specific endpoint when creating a new client. * Added a mocked instance admin server for testing purpose. * Added an environment variable "GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING" to enable this feature. By default, it is disabled. * If there is a PermissionDenied error when calling GetInstance(), fallback to use the global endpoint or the user-specified endpoint. * Included tests to verify the functionality. Change-Id: Ie447d13b8e414f6299fc56acb678d293531849ae Reviewed-on: https://code-review.googlesource.com/c/gocloud/+/48895 Reviewed-by: kokoro Reviewed-by: Reid Hironaga Reviewed-by: Shanika Kuruppu --- spanner/client.go | 78 +++++++- spanner/client_test.go | 187 +++++++++++++++++- spanner/integration_test.go | 46 ++++- .../testutil/inmem_instance_admin_server.go | 86 ++++++++ .../inmem_instance_admin_server_test.go | 95 +++++++++ .../internal/testutil/mocked_inmem_server.go | 27 ++- 6 files changed, 509 insertions(+), 10 deletions(-) create mode 100644 spanner/internal/testutil/inmem_instance_admin_server.go create mode 100644 spanner/internal/testutil/inmem_instance_admin_server_test.go diff --git a/spanner/client.go b/spanner/client.go index ae05960b953e..4f0ee8e7a5d3 100644 --- a/spanner/client.go +++ b/spanner/client.go @@ -25,12 +25,16 @@ import ( "time" "cloud.google.com/go/internal/trace" + instance "cloud.google.com/go/spanner/admin/instance/apiv1" vkit "cloud.google.com/go/spanner/apiv1" "google.golang.org/api/option" + instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" sppb "google.golang.org/genproto/googleapis/spanner/v1" + field_mask "google.golang.org/genproto/protobuf/field_mask" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" ) const ( @@ -50,7 +54,8 @@ const ( ) var ( - validDBPattern = regexp.MustCompile("^projects/[^/]+/instances/[^/]+/databases/[^/]+$") + validDBPattern = regexp.MustCompile("^projects/[^/]+/instances/[^/]+/databases/[^/]+$") + validInstancePattern = regexp.MustCompile("^projects/[^/]+/instances/[^/]+") ) func validDatabaseName(db string) error { @@ -61,6 +66,15 @@ func validDatabaseName(db string) error { return nil } +func getInstanceName(db string) (string, error) { + matches := validInstancePattern.FindStringSubmatch(db) + if len(matches) == 0 { + return "", fmt.Errorf("Failed to retrieve instance name from %q according to pattern %q", + db, validInstancePattern.String()) + } + return matches[0], nil +} + // Client is a client for reading and writing data to a Cloud Spanner database. // A client is safe to use concurrently, except for its Close method. type Client struct { @@ -103,6 +117,42 @@ func contextWithOutgoingMetadata(ctx context.Context, md metadata.MD) context.Co return metadata.NewOutgoingContext(ctx, md) } +// getInstanceEndpoint returns an instance-specific endpoint if one exists. If +// multiple endpoints exist, it returns the first one. +func getInstanceEndpoint(ctx context.Context, database string, opts ...option.ClientOption) (string, error) { + instanceName, err := getInstanceName(database) + if err != nil { + return "", fmt.Errorf("Failed to resolve endpoint: %v", err) + } + + c, err := instance.NewInstanceAdminClient(ctx, opts...) + if err != nil { + return "", err + } + defer c.Close() + + req := &instancepb.GetInstanceRequest{ + Name: instanceName, + FieldMask: &field_mask.FieldMask{ + Paths: []string{"endpoint_uris"}, + }, + } + + resp, err := c.GetInstance(ctx, req) + if err != nil { + return "", err + } + + endpointURIs := resp.GetEndpointUris() + + if len(endpointURIs) > 0 { + return endpointURIs[0], nil + } + + // Return empty string when no endpoints exist. + return "", nil +} + // NewClient creates a client to a database. A valid database name has the // form projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID. It uses // a default configuration. @@ -142,6 +192,32 @@ func NewClientWithConfig(ctx context.Context, database string, config ClientConf option.WithoutAuthentication(), } opts = append(opts, emulatorOpts...) + } else if os.Getenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING") == "true" { + // Fetch the instance-specific endpoint. + reqOpts := []option.ClientOption{option.WithEndpoint(endpoint)} + reqOpts = append(reqOpts, opts...) + instanceEndpoint, err := getInstanceEndpoint(ctx, database, reqOpts...) + + if err != nil { + // If there is a PermissionDenied error, fall back to use the global endpoint + // or the user-specified endpoint. + if status.Code(err) == codes.PermissionDenied { + logf(config.logger, ` +Warning: The client library attempted to connect to an endpoint closer to your +Cloud Spanner data but was unable to do so. The client library will fall back +and route requests to the global Spanner endpoint (spanner.googleapis.com), +which may result in increased latency. We recommend including the scope +https://www.googleapis.com/auth/spanner.admin so that the client library can +get an instance-specific endpoint and efficiently route requests. +`) + } else { + return nil, err + } + } + + if instanceEndpoint != "" { + opts = append(opts, option.WithEndpoint(instanceEndpoint)) + } } // gRPC options. diff --git a/spanner/client_test.go b/spanner/client_test.go index 8f08f38c3fc5..40cf693251c6 100644 --- a/spanner/client_test.go +++ b/spanner/client_test.go @@ -20,14 +20,20 @@ import ( "context" "fmt" "io" + "io/ioutil" + "log" + "os" "strings" "testing" "time" itestutil "cloud.google.com/go/internal/testutil" . "cloud.google.com/go/spanner/internal/testutil" + "github.com/golang/protobuf/proto" "google.golang.org/api/iterator" "google.golang.org/api/option" + instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" + sppb "google.golang.org/genproto/googleapis/spanner/v1" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) @@ -94,6 +100,34 @@ func TestValidDatabaseName(t *testing.T) { } } +// Test getInstanceName() +func TestGetInstanceName(t *testing.T) { + validDbURI := "projects/spanner-cloud-test/instances/foo/databases/foodb" + invalidDbUris := []string{ + // Completely wrong DB URI. + "foobarDB", + // Project ID contains "/". + "projects/spanner-cloud/test/instances/foo/databases/foodb", + // No instance ID. + "projects/spanner-cloud-test/instances//databases/foodb", + } + want := "projects/spanner-cloud-test/instances/foo" + got, err := getInstanceName(validDbURI) + if err != nil { + t.Errorf("getInstanceName(%q) has an error: %q, want nil", validDbURI, err) + } + if got != want { + t.Errorf("getInstanceName(%q) = %q, want %q", validDbURI, got, want) + } + for _, d := range invalidDbUris { + wantErr := "Failed to retrieve instance name" + _, err = getInstanceName(d) + if !strings.Contains(err.Error(), wantErr) { + t.Errorf("getInstanceName(%q) has an error: %q, want error pattern %q", validDbURI, err, wantErr) + } + } +} + func TestReadOnlyTransactionClose(t *testing.T) { // Closing a ReadOnlyTransaction shouldn't panic. c := &Client{} @@ -121,7 +155,7 @@ func TestClient_Single_InvalidArgument(t *testing.T) { t.Parallel() err := testSingleQuery(t, status.Error(codes.InvalidArgument, "Invalid argument")) if status.Code(err) != codes.InvalidArgument { - t.Fatalf("got unexpected exception %v, expected InvalidArgument", err) + t.Fatalf("got: %v, want: %v", err, codes.InvalidArgument) } } @@ -289,6 +323,157 @@ func TestClient_Single_ContextCanceled_withDeclaredServerErrors(t *testing.T) { } } +func TestClient_ResourceBasedRouting_WithEndpointsReturned(t *testing.T) { + os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "true") + defer os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "") + + // Create two servers. The base server receives the GetInstance request and + // returns the instance endpoint of the target server. The client should contact + // the target server after getting the instance endpoint. + serverBase, optsBase, serverTeardownBase := NewMockedSpannerInMemTestServerWithAddr(t, "localhost:8081") + defer serverTeardownBase() + serverTarget, optsTarget, serverTeardownTarget := NewMockedSpannerInMemTestServerWithAddr(t, "localhost:8082") + defer serverTeardownTarget() + + // Return the instance endpoint. + instanceEndpoint := fmt.Sprintf("%s", optsTarget[0]) + resps := []proto.Message{&instancepb.Instance{ + EndpointUris: []string{instanceEndpoint}, + }} + serverBase.TestInstanceAdmin.SetResps(resps) + + ctx := context.Background() + formattedDatabase := fmt.Sprintf("projects/%s/instances/%s/databases/%s", "some-project", "some-instance", "some-database") + client, err := NewClientWithConfig(ctx, formattedDatabase, ClientConfig{}, optsBase...) + if err != nil { + t.Fatal(err) + } + + if err := executeSingerQuery(ctx, client.Single()); err != nil { + t.Fatal(err) + } + + // The base server should not receive any requests. + if _, err := shouldHaveReceived(serverBase.TestSpanner, []interface{}{}); err != nil { + t.Fatal(err) + } + + // The target server should receive requests. + if _, err = shouldHaveReceived(serverTarget.TestSpanner, []interface{}{ + &sppb.CreateSessionRequest{}, + &sppb.ExecuteSqlRequest{}, + }); err != nil { + t.Fatal(err) + } +} + +func TestClient_ResourceBasedRouting_WithoutEndpointsReturned(t *testing.T) { + os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "true") + defer os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "") + + server, opts, serverTeardown := NewMockedSpannerInMemTestServer(t) + defer serverTeardown() + + // Return an empty list of endpoints. + resps := []proto.Message{&instancepb.Instance{ + EndpointUris: []string{}, + }} + server.TestInstanceAdmin.SetResps(resps) + + ctx := context.Background() + formattedDatabase := fmt.Sprintf("projects/%s/instances/%s/databases/%s", "some-project", "some-instance", "some-database") + client, err := NewClientWithConfig(ctx, formattedDatabase, ClientConfig{}, opts...) + if err != nil { + t.Fatal(err) + } + + if err := executeSingerQuery(ctx, client.Single()); err != nil { + t.Fatal(err) + } + + // Check if the request goes to the default endpoint. + if _, err := shouldHaveReceived(server.TestSpanner, []interface{}{ + &sppb.CreateSessionRequest{}, + &sppb.ExecuteSqlRequest{}, + }); err != nil { + t.Fatal(err) + } +} + +func TestClient_ResourceBasedRouting_WithPermissionDeniedError(t *testing.T) { + os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "true") + defer os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "") + + server, opts, serverTeardown := NewMockedSpannerInMemTestServer(t) + defer serverTeardown() + + server.TestInstanceAdmin.SetErr(status.Error(codes.PermissionDenied, "Permission Denied")) + + ctx := context.Background() + formattedDatabase := fmt.Sprintf("projects/%s/instances/%s/databases/%s", "some-project", "some-instance", "some-database") + // `PermissionDeniedError` causes a warning message to be logged, which is expected. + // We set the output to be discarded to avoid spamming the log. + logger := log.New(ioutil.Discard, "", log.LstdFlags) + client, err := NewClientWithConfig(ctx, formattedDatabase, ClientConfig{logger: logger}, opts...) + if err != nil { + t.Fatal(err) + } + + if err := executeSingerQuery(ctx, client.Single()); err != nil { + t.Fatal(err) + } + + // Fallback to use the default endpoint when calling GetInstance() returns + // a PermissionDenied error. + if _, err := shouldHaveReceived(server.TestSpanner, []interface{}{ + &sppb.CreateSessionRequest{}, + &sppb.ExecuteSqlRequest{}, + }); err != nil { + t.Fatal(err) + } +} + +func TestClient_ResourceBasedRouting_WithUnavailableError(t *testing.T) { + os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "true") + defer os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "") + + server, opts, serverTeardown := NewMockedSpannerInMemTestServer(t) + defer serverTeardown() + + resps := []proto.Message{&instancepb.Instance{ + EndpointUris: []string{}, + }} + server.TestInstanceAdmin.SetResps(resps) + server.TestInstanceAdmin.SetErr(status.Error(codes.Unavailable, "Temporary unavailable")) + + ctx := context.Background() + formattedDatabase := fmt.Sprintf("projects/%s/instances/%s/databases/%s", "some-project", "some-instance", "some-database") + _, err := NewClientWithConfig(ctx, formattedDatabase, ClientConfig{}, opts...) + // The first request will get an error and the server resets the error to nil, + // so the next request will be fine. Due to retrying, there is no errors. + if err != nil { + t.Fatal(err) + } +} + +func TestClient_ResourceBasedRouting_WithInvalidArgumentError(t *testing.T) { + os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "true") + defer os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "") + + server, opts, serverTeardown := NewMockedSpannerInMemTestServer(t) + defer serverTeardown() + + server.TestInstanceAdmin.SetErr(status.Error(codes.InvalidArgument, "Invalid argument")) + + ctx := context.Background() + formattedDatabase := fmt.Sprintf("projects/%s/instances/%s/databases/%s", "some-project", "some-instance", "some-database") + _, err := NewClientWithConfig(ctx, formattedDatabase, ClientConfig{}, opts...) + + if status.Code(err) != codes.InvalidArgument { + t.Fatalf("got unexpected exception %v, expected InvalidArgument", err) + } +} + func testSingleQuery(t *testing.T, serverError error) error { ctx := context.Background() server, client, teardown := setupMockedTestServer(t) diff --git a/spanner/integration_test.go b/spanner/integration_test.go index 70c19c96c23a..92324d3b37b3 100644 --- a/spanner/integration_test.go +++ b/spanner/integration_test.go @@ -440,6 +440,50 @@ func TestIntegration_SingleUse(t *testing.T) { } } +// Test resource-based routing enabled. +func TestIntegration_SingleUse_WithResourceBasedRouting(t *testing.T) { + os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "true") + defer os.Setenv("GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING", "") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + // Set up testing environment. + client, _, cleanup := prepareIntegrationTest(ctx, t, DefaultSessionPoolConfig, singerDBStatements) + defer cleanup() + + writes := []struct { + row []interface{} + ts time.Time + }{ + {row: []interface{}{1, "Marc", "Foo"}}, + {row: []interface{}{2, "Tars", "Bar"}}, + {row: []interface{}{3, "Alpha", "Beta"}}, + {row: []interface{}{4, "Last", "End"}}, + } + // Try to write four rows through the Apply API. + for i, w := range writes { + var err error + m := InsertOrUpdate("Singers", + []string{"SingerId", "FirstName", "LastName"}, + w.row) + if writes[i].ts, err = client.Apply(ctx, []*Mutation{m}, ApplyAtLeastOnce()); err != nil { + t.Fatal(err) + } + } + + row, err := client.Single().ReadRow(ctx, "Singers", Key{3}, []string{"FirstName"}) + if err != nil { + t.Errorf("SingleUse.ReadRow returns error %v, want nil", err) + } + var got string + if err := row.Column(0, &got); err != nil { + t.Errorf("row.Column returns error %v, want nil", err) + } + if want := "Alpha"; got != want { + t.Errorf("got %q, want %q", got, want) + } +} + func TestIntegration_SingleUse_ReadingWithLimit(t *testing.T) { t.Parallel() @@ -2594,7 +2638,7 @@ func isNaN(x interface{}) bool { func createClient(ctx context.Context, dbPath string, spc SessionPoolConfig) (client *Client, err error) { client, err = NewClientWithConfig(ctx, dbPath, ClientConfig{ SessionPoolConfig: spc, - }, option.WithTokenSource(testutil.TokenSource(ctx, Scope)), option.WithEndpoint(endpoint)) + }, option.WithTokenSource(testutil.TokenSource(ctx, Scope, AdminScope)), option.WithEndpoint(endpoint)) if err != nil { return nil, fmt.Errorf("cannot create data client on DB %v: %v", dbPath, err) } diff --git a/spanner/internal/testutil/inmem_instance_admin_server.go b/spanner/internal/testutil/inmem_instance_admin_server.go new file mode 100644 index 000000000000..f02dce0a442a --- /dev/null +++ b/spanner/internal/testutil/inmem_instance_admin_server.go @@ -0,0 +1,86 @@ +// Copyright 2020 Google LLC +// +// 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 +// +// https://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 testutil + +import ( + "context" + + "github.com/golang/protobuf/proto" + instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" +) + +// InMemInstanceAdminServer contains the InstanceAdminServer interface plus a couple +// of specific methods for setting mocked results. +type InMemInstanceAdminServer interface { + instancepb.InstanceAdminServer + Stop() + Resps() []proto.Message + SetResps([]proto.Message) + Reqs() []proto.Message + SetReqs([]proto.Message) + SetErr(error) +} + +// inMemInstanceAdminServer implements InMemInstanceAdminServer interface. Note that +// there is no mutex protecting the data structures, so it is not safe for +// concurrent use. +type inMemInstanceAdminServer struct { + instancepb.InstanceAdminServer + reqs []proto.Message + // If set, all calls return this error + err error + // responses to return if err == nil + resps []proto.Message +} + +// NewInMemInstanceAdminServer creates a new in-mem test server. +func NewInMemInstanceAdminServer() InMemInstanceAdminServer { + res := &inMemInstanceAdminServer{} + return res +} + +// GetInstance returns the metadata of a spanner instance. +func (s *inMemInstanceAdminServer) GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest) (*instancepb.Instance, error) { + s.reqs = append(s.reqs, req) + if s.err != nil { + defer func() { s.err = nil }() + return nil, s.err + } + return s.resps[0].(*instancepb.Instance), nil +} + +func (s *inMemInstanceAdminServer) Stop() { + // do nothing +} + +func (s *inMemInstanceAdminServer) Resps() []proto.Message { + return s.resps +} + +func (s *inMemInstanceAdminServer) SetResps(resps []proto.Message) { + s.resps = resps +} + +func (s *inMemInstanceAdminServer) Reqs() []proto.Message { + return s.reqs +} + +func (s *inMemInstanceAdminServer) SetReqs(reqs []proto.Message) { + s.reqs = reqs +} + +func (s *inMemInstanceAdminServer) SetErr(err error) { + s.err = err +} diff --git a/spanner/internal/testutil/inmem_instance_admin_server_test.go b/spanner/internal/testutil/inmem_instance_admin_server_test.go new file mode 100644 index 000000000000..f4e3b22b7202 --- /dev/null +++ b/spanner/internal/testutil/inmem_instance_admin_server_test.go @@ -0,0 +1,95 @@ +// Copyright 2020 Google LLC +// +// 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 +// +// https://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 testutil_test + +import ( + "context" + "flag" + "fmt" + "log" + "net" + "testing" + + instance "cloud.google.com/go/spanner/admin/instance/apiv1" + "cloud.google.com/go/spanner/internal/testutil" + "github.com/golang/protobuf/proto" + "google.golang.org/api/option" + instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" + "google.golang.org/grpc" +) + +var instanceClientOpt option.ClientOption + +var ( + mockInstanceAdmin = testutil.NewInMemInstanceAdminServer() +) + +func setupInstanceAdminServer() { + flag.Parse() + + serv := grpc.NewServer() + instancepb.RegisterInstanceAdminServer(serv, mockInstanceAdmin) + + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + log.Fatal(err) + } + go serv.Serve(lis) + + conn, err := grpc.Dial(lis.Addr().String(), grpc.WithInsecure()) + if err != nil { + log.Fatal(err) + } + instanceClientOpt = option.WithGRPCConn(conn) +} + +func TestInstanceAdminGetInstance(t *testing.T) { + setupInstanceAdminServer() + var expectedResponse = &instancepb.Instance{ + Name: "name2-1052831874", + Config: "config-1354792126", + DisplayName: "displayName1615086568", + NodeCount: 1539922066, + } + + mockInstanceAdmin.SetErr(nil) + mockInstanceAdmin.SetReqs(nil) + + mockInstanceAdmin.SetResps(append(mockInstanceAdmin.Resps()[:0], expectedResponse)) + + var formattedName string = fmt.Sprintf("projects/%s/instances/%s", "[PROJECT]", "[INSTANCE]") + var request = &instancepb.GetInstanceRequest{ + Name: formattedName, + } + + c, err := instance.NewInstanceAdminClient(context.Background(), instanceClientOpt) + if err != nil { + t.Fatal(err) + } + + resp, err := c.GetInstance(context.Background(), request) + + if err != nil { + t.Fatal(err) + } + + if want, got := request, mockInstanceAdmin.Reqs()[0]; !proto.Equal(want, got) { + t.Errorf("wrong request %q, want %q", got, want) + } + + if want, got := expectedResponse, resp; !proto.Equal(want, got) { + t.Errorf("wrong response %q, want %q)", got, want) + } +} diff --git a/spanner/internal/testutil/mocked_inmem_server.go b/spanner/internal/testutil/mocked_inmem_server.go index 9ac328fca510..028f038eb761 100644 --- a/spanner/internal/testutil/mocked_inmem_server.go +++ b/spanner/internal/testutil/mocked_inmem_server.go @@ -22,6 +22,7 @@ import ( structpb "github.com/golang/protobuf/ptypes/struct" "google.golang.org/api/option" + instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" spannerpb "google.golang.org/genproto/googleapis/spanner/v1" "google.golang.org/grpc" ) @@ -58,29 +59,41 @@ const UpdateBarSetFooRowCount = 5 // MockedSpannerInMemTestServer is an InMemSpannerServer with results for a // number of SQL statements readily mocked. type MockedSpannerInMemTestServer struct { - TestSpanner InMemSpannerServer - server *grpc.Server + TestSpanner InMemSpannerServer + TestInstanceAdmin InMemInstanceAdminServer + server *grpc.Server } -// NewMockedSpannerInMemTestServer creates a MockedSpannerInMemTestServer and -// returns client options that can be used to connect to it. +// NewMockedSpannerInMemTestServer creates a MockedSpannerInMemTestServer at +// localhost with a random port and returns client options that can be used +// to connect to it. func NewMockedSpannerInMemTestServer(t *testing.T) (mockedServer *MockedSpannerInMemTestServer, opts []option.ClientOption, teardown func()) { + return NewMockedSpannerInMemTestServerWithAddr(t, "localhost:0") +} + +// NewMockedSpannerInMemTestServerWithAddr creates a MockedSpannerInMemTestServer +// at a given listening address and returns client options that can be used +// to connect to it. +func NewMockedSpannerInMemTestServerWithAddr(t *testing.T, addr string) (mockedServer *MockedSpannerInMemTestServer, opts []option.ClientOption, teardown func()) { mockedServer = &MockedSpannerInMemTestServer{} - opts = mockedServer.setupMockedServer(t) + opts = mockedServer.setupMockedServerWithAddr(t, addr) return mockedServer, opts, func() { mockedServer.TestSpanner.Stop() + mockedServer.TestInstanceAdmin.Stop() mockedServer.server.Stop() } } -func (s *MockedSpannerInMemTestServer) setupMockedServer(t *testing.T) []option.ClientOption { +func (s *MockedSpannerInMemTestServer) setupMockedServerWithAddr(t *testing.T, addr string) []option.ClientOption { s.TestSpanner = NewInMemSpannerServer() + s.TestInstanceAdmin = NewInMemInstanceAdminServer() s.setupFooResults() s.setupSingersResults() s.server = grpc.NewServer() spannerpb.RegisterSpannerServer(s.server, s.TestSpanner) + instancepb.RegisterInstanceAdminServer(s.server, s.TestInstanceAdmin) - lis, err := net.Listen("tcp", "localhost:0") + lis, err := net.Listen("tcp", addr) if err != nil { t.Fatal(err) }