1// Copyright 2019 Google LLC 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// https://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package testutil 16 17import ( 18 "fmt" 19 "net" 20 "strconv" 21 "testing" 22 23 structpb "github.com/golang/protobuf/ptypes/struct" 24 "google.golang.org/api/option" 25 instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" 26 spannerpb "google.golang.org/genproto/googleapis/spanner/v1" 27 "google.golang.org/grpc" 28) 29 30// SelectFooFromBar is a SELECT statement that is added to the mocked test 31// server and will return a one-col-two-rows result set containing the INT64 32// values 1 and 2. 33const SelectFooFromBar = "SELECT FOO FROM BAR" 34const selectFooFromBarRowCount int64 = 2 35const selectFooFromBarColCount int = 1 36 37var selectFooFromBarResults = [...]int64{1, 2} 38 39// SelectSingerIDAlbumIDAlbumTitleFromAlbums i a SELECT statement that is added 40// to the mocked test server and will return a 3-cols-3-rows result set. 41const SelectSingerIDAlbumIDAlbumTitleFromAlbums = "SELECT SingerId, AlbumId, AlbumTitle FROM Albums" 42 43// SelectSingerIDAlbumIDAlbumTitleFromAlbumsRowCount is the number of rows 44// returned by the SelectSingerIDAlbumIDAlbumTitleFromAlbums statement. 45const SelectSingerIDAlbumIDAlbumTitleFromAlbumsRowCount int64 = 3 46 47// SelectSingerIDAlbumIDAlbumTitleFromAlbumsColCount is the number of cols 48// returned by the SelectSingerIDAlbumIDAlbumTitleFromAlbums statement. 49const SelectSingerIDAlbumIDAlbumTitleFromAlbumsColCount int = 3 50 51// UpdateBarSetFoo is an UPDATE statement that is added to the mocked test 52// server that will return an update count of 5. 53const UpdateBarSetFoo = "UPDATE FOO SET BAR=1 WHERE BAZ=2" 54 55// UpdateBarSetFooRowCount is the constant update count value returned by the 56// statement defined in UpdateBarSetFoo. 57const UpdateBarSetFooRowCount = 5 58 59// MockedSpannerInMemTestServer is an InMemSpannerServer with results for a 60// number of SQL statements readily mocked. 61type MockedSpannerInMemTestServer struct { 62 TestSpanner InMemSpannerServer 63 TestInstanceAdmin InMemInstanceAdminServer 64 server *grpc.Server 65} 66 67// NewMockedSpannerInMemTestServer creates a MockedSpannerInMemTestServer at 68// localhost with a random port and returns client options that can be used 69// to connect to it. 70func NewMockedSpannerInMemTestServer(t *testing.T) (mockedServer *MockedSpannerInMemTestServer, opts []option.ClientOption, teardown func()) { 71 return NewMockedSpannerInMemTestServerWithAddr(t, "localhost:0") 72} 73 74// NewMockedSpannerInMemTestServerWithAddr creates a MockedSpannerInMemTestServer 75// at a given listening address and returns client options that can be used 76// to connect to it. 77func NewMockedSpannerInMemTestServerWithAddr(t *testing.T, addr string) (mockedServer *MockedSpannerInMemTestServer, opts []option.ClientOption, teardown func()) { 78 mockedServer = &MockedSpannerInMemTestServer{} 79 opts = mockedServer.setupMockedServerWithAddr(t, addr) 80 return mockedServer, opts, func() { 81 mockedServer.TestSpanner.Stop() 82 mockedServer.TestInstanceAdmin.Stop() 83 mockedServer.server.Stop() 84 } 85} 86 87func (s *MockedSpannerInMemTestServer) setupMockedServerWithAddr(t *testing.T, addr string) []option.ClientOption { 88 s.TestSpanner = NewInMemSpannerServer() 89 s.TestInstanceAdmin = NewInMemInstanceAdminServer() 90 s.setupFooResults() 91 s.setupSingersResults() 92 s.server = grpc.NewServer() 93 spannerpb.RegisterSpannerServer(s.server, s.TestSpanner) 94 instancepb.RegisterInstanceAdminServer(s.server, s.TestInstanceAdmin) 95 96 lis, err := net.Listen("tcp", addr) 97 if err != nil { 98 t.Fatal(err) 99 } 100 go s.server.Serve(lis) 101 102 serverAddress := lis.Addr().String() 103 opts := []option.ClientOption{ 104 option.WithEndpoint(serverAddress), 105 option.WithGRPCDialOption(grpc.WithInsecure()), 106 option.WithoutAuthentication(), 107 } 108 return opts 109} 110 111func (s *MockedSpannerInMemTestServer) setupFooResults() { 112 fields := make([]*spannerpb.StructType_Field, selectFooFromBarColCount) 113 fields[0] = &spannerpb.StructType_Field{ 114 Name: "FOO", 115 Type: &spannerpb.Type{Code: spannerpb.TypeCode_INT64}, 116 } 117 rowType := &spannerpb.StructType{ 118 Fields: fields, 119 } 120 metadata := &spannerpb.ResultSetMetadata{ 121 RowType: rowType, 122 } 123 rows := make([]*structpb.ListValue, selectFooFromBarRowCount) 124 for idx, value := range selectFooFromBarResults { 125 rowValue := make([]*structpb.Value, selectFooFromBarColCount) 126 rowValue[0] = &structpb.Value{ 127 Kind: &structpb.Value_StringValue{StringValue: strconv.FormatInt(value, 10)}, 128 } 129 rows[idx] = &structpb.ListValue{ 130 Values: rowValue, 131 } 132 } 133 resultSet := &spannerpb.ResultSet{ 134 Metadata: metadata, 135 Rows: rows, 136 } 137 result := &StatementResult{Type: StatementResultResultSet, ResultSet: resultSet} 138 s.TestSpanner.PutStatementResult(SelectFooFromBar, result) 139 s.TestSpanner.PutStatementResult(UpdateBarSetFoo, &StatementResult{ 140 Type: StatementResultUpdateCount, 141 UpdateCount: UpdateBarSetFooRowCount, 142 }) 143} 144 145func (s *MockedSpannerInMemTestServer) setupSingersResults() { 146 fields := make([]*spannerpb.StructType_Field, SelectSingerIDAlbumIDAlbumTitleFromAlbumsColCount) 147 fields[0] = &spannerpb.StructType_Field{ 148 Name: "SingerId", 149 Type: &spannerpb.Type{Code: spannerpb.TypeCode_INT64}, 150 } 151 fields[1] = &spannerpb.StructType_Field{ 152 Name: "AlbumId", 153 Type: &spannerpb.Type{Code: spannerpb.TypeCode_INT64}, 154 } 155 fields[2] = &spannerpb.StructType_Field{ 156 Name: "AlbumTitle", 157 Type: &spannerpb.Type{Code: spannerpb.TypeCode_STRING}, 158 } 159 rowType := &spannerpb.StructType{ 160 Fields: fields, 161 } 162 metadata := &spannerpb.ResultSetMetadata{ 163 RowType: rowType, 164 } 165 rows := make([]*structpb.ListValue, SelectSingerIDAlbumIDAlbumTitleFromAlbumsRowCount) 166 var idx int64 167 for idx = 0; idx < SelectSingerIDAlbumIDAlbumTitleFromAlbumsRowCount; idx++ { 168 rowValue := make([]*structpb.Value, SelectSingerIDAlbumIDAlbumTitleFromAlbumsColCount) 169 rowValue[0] = &structpb.Value{ 170 Kind: &structpb.Value_StringValue{StringValue: strconv.FormatInt(idx+1, 10)}, 171 } 172 rowValue[1] = &structpb.Value{ 173 Kind: &structpb.Value_StringValue{StringValue: strconv.FormatInt(idx*10+idx, 10)}, 174 } 175 rowValue[2] = &structpb.Value{ 176 Kind: &structpb.Value_StringValue{StringValue: fmt.Sprintf("Album title %d", idx)}, 177 } 178 rows[idx] = &structpb.ListValue{ 179 Values: rowValue, 180 } 181 } 182 resultSet := &spannerpb.ResultSet{ 183 Metadata: metadata, 184 Rows: rows, 185 } 186 result := &StatementResult{Type: StatementResultResultSet, ResultSet: resultSet} 187 s.TestSpanner.PutStatementResult(SelectSingerIDAlbumIDAlbumTitleFromAlbums, result) 188} 189