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