spanner/internal/testutil/mockclient.go

/*
Copyright 2017 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

    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 testutil

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes/empty"
	proto3 "github.com/golang/protobuf/ptypes/struct"
	pbt "github.com/golang/protobuf/ptypes/timestamp"
	pbs "google.golang.org/genproto/googleapis/rpc/status"
	sppb "google.golang.org/genproto/googleapis/spanner/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// MockCloudSpannerClient is a mock implementation of sppb.SpannerClient.
type MockCloudSpannerClient struct {
	sppb.SpannerClient

	mu sync.Mutex
	t  *testing.T
	// Live sessions on the client.
	sessions map[string]bool
	// Session ping history.
	pings []string
	// Client will stall on any requests.
	freezed chan struct{}

	// Expected set of actions that have been executed by the client. These
	// interfaces should be type reflected against with *Request types in sppb,
	// such as sppb.GetSessionRequest. Buffered to a large degree.
	ReceivedRequests chan interface{}
}

// NewMockCloudSpannerClient creates new MockCloudSpannerClient instance.
func NewMockCloudSpannerClient(t *testing.T) *MockCloudSpannerClient {
	mc := &MockCloudSpannerClient{
		t:                t,
		sessions:         map[string]bool{},
		ReceivedRequests: make(chan interface{}, 100000),
	}

	// Produce a closed channel, so the default action of ready is to not block.
	mc.Freeze()
	mc.Unfreeze()

	return mc
}

// DumpPings dumps the ping history.
func (m *MockCloudSpannerClient) DumpPings() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	return append([]string(nil), m.pings...)
}

// DumpSessions dumps the internal session table.
func (m *MockCloudSpannerClient) DumpSessions() map[string]bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	st := map[string]bool{}
	for s, v := range m.sessions {
		st[s] = v
	}
	return st
}

// CreateSession is a placeholder for SpannerClient.CreateSession.
func (m *MockCloudSpannerClient) CreateSession(ctx context.Context, r *sppb.CreateSessionRequest, opts ...grpc.CallOption) (*sppb.Session, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	s := &sppb.Session{}
	if r.Database != "mockdb" {
		// Reject other databases
		return s, status.Errorf(codes.NotFound, fmt.Sprintf("database not found: %v", r.Database))
	}
	// Generate & record session name.
	s.Name = fmt.Sprintf("mockdb-%v", time.Now().UnixNano())
	m.sessions[s.Name] = true
	return s, nil
}

// GetSession is a placeholder for SpannerClient.GetSession.
func (m *MockCloudSpannerClient) GetSession(ctx context.Context, r *sppb.GetSessionRequest, opts ...grpc.CallOption) (*sppb.Session, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	m.pings = append(m.pings, r.Name)
	if _, ok := m.sessions[r.Name]; !ok {
		return nil, status.Errorf(codes.NotFound, fmt.Sprintf("Session not found: %v", r.Name))
	}
	return &sppb.Session{Name: r.Name}, nil
}

// DeleteSession is a placeholder for SpannerClient.DeleteSession.
func (m *MockCloudSpannerClient) DeleteSession(ctx context.Context, r *sppb.DeleteSessionRequest, opts ...grpc.CallOption) (*empty.Empty, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.sessions[r.Name]; !ok {
		// Session not found.
		return &empty.Empty{}, status.Errorf(codes.NotFound, fmt.Sprintf("Session not found: %v", r.Name))
	}
	// Delete session from in-memory table.
	delete(m.sessions, r.Name)
	return &empty.Empty{}, nil
}

// ExecuteSql is a placeholder for SpannerClient.ExecuteSql.
func (m *MockCloudSpannerClient) ExecuteSql(ctx context.Context, r *sppb.ExecuteSqlRequest, opts ...grpc.CallOption) (*sppb.ResultSet, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	return &sppb.ResultSet{Stats: &sppb.ResultSetStats{RowCount: &sppb.ResultSetStats_RowCountExact{7}}}, nil
}

// ExecuteBatchDml is a placeholder for SpannerClient.ExecuteBatchDml.
func (m *MockCloudSpannerClient) ExecuteBatchDml(ctx context.Context, r *sppb.ExecuteBatchDmlRequest, opts ...grpc.CallOption) (*sppb.ExecuteBatchDmlResponse, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	return &sppb.ExecuteBatchDmlResponse{Status: &pbs.Status{Code: 0}, ResultSets: []*sppb.ResultSet{}}, nil
}

// ExecuteStreamingSql is a mock implementation of SpannerClient.ExecuteStreamingSql.
func (m *MockCloudSpannerClient) ExecuteStreamingSql(ctx context.Context, r *sppb.ExecuteSqlRequest, opts ...grpc.CallOption) (sppb.Spanner_ExecuteStreamingSqlClient, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	wantReq := &sppb.ExecuteSqlRequest{
		Session: "mocksession",
		Transaction: &sppb.TransactionSelector{
			Selector: &sppb.TransactionSelector_SingleUse{
				SingleUse: &sppb.TransactionOptions{
					Mode: &sppb.TransactionOptions_ReadOnly_{
						ReadOnly: &sppb.TransactionOptions_ReadOnly{
							TimestampBound: &sppb.TransactionOptions_ReadOnly_Strong{
								Strong: true,
							},
							ReturnReadTimestamp: false,
						},
					},
				},
			},
		},
		Sql: "mockquery",
		Params: &proto3.Struct{
			Fields: map[string]*proto3.Value{"var1": {Kind: &proto3.Value_StringValue{StringValue: "abc"}}},
		},
		ParamTypes: map[string]*sppb.Type{"var1": {Code: sppb.TypeCode_STRING}},
	}
	if !proto.Equal(r, wantReq) {
		return nil, fmt.Errorf("got query request: %v, want: %v", r, wantReq)
	}
	return nil, errors.New("query never succeeds on mock client")
}

// StreamingRead is a placeholder for SpannerClient.StreamingRead.
func (m *MockCloudSpannerClient) StreamingRead(ctx context.Context, r *sppb.ReadRequest, opts ...grpc.CallOption) (sppb.Spanner_StreamingReadClient, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	wantReq := &sppb.ReadRequest{
		Session: "mocksession",
		Transaction: &sppb.TransactionSelector{
			Selector: &sppb.TransactionSelector_SingleUse{
				SingleUse: &sppb.TransactionOptions{
					Mode: &sppb.TransactionOptions_ReadOnly_{
						ReadOnly: &sppb.TransactionOptions_ReadOnly{
							TimestampBound: &sppb.TransactionOptions_ReadOnly_Strong{
								Strong: true,
							},
							ReturnReadTimestamp: false,
						},
					},
				},
			},
		},
		Table:   "t_mock",
		Columns: []string{"col1", "col2"},
		KeySet: &sppb.KeySet{
			Keys: []*proto3.ListValue{
				{
					Values: []*proto3.Value{
						{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
					},
				},
			},
			Ranges: []*sppb.KeyRange{},
			All:    false,
		},
	}
	if !proto.Equal(r, wantReq) {
		return nil, fmt.Errorf("got query request: %v, want: %v", r, wantReq)
	}
	return nil, errors.New("read never succeeds on mock client")
}

// BeginTransaction is a placeholder for SpannerClient.BeginTransaction.
func (m *MockCloudSpannerClient) BeginTransaction(ctx context.Context, r *sppb.BeginTransactionRequest, opts ...grpc.CallOption) (*sppb.Transaction, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	resp := &sppb.Transaction{Id: []byte("transaction-1")}
	if _, ok := r.Options.Mode.(*sppb.TransactionOptions_ReadOnly_); ok {
		resp.ReadTimestamp = &pbt.Timestamp{Seconds: 3, Nanos: 4}
	}
	return resp, nil
}

// Commit is a placeholder for SpannerClient.Commit.
func (m *MockCloudSpannerClient) Commit(ctx context.Context, r *sppb.CommitRequest, opts ...grpc.CallOption) (*sppb.CommitResponse, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	return &sppb.CommitResponse{CommitTimestamp: &pbt.Timestamp{Seconds: 1, Nanos: 2}}, nil
}

// Rollback is a placeholder for SpannerClient.Rollback.
func (m *MockCloudSpannerClient) Rollback(ctx context.Context, r *sppb.RollbackRequest, opts ...grpc.CallOption) (*empty.Empty, error) {
	m.ready()
	m.ReceivedRequests <- r

	m.mu.Lock()
	defer m.mu.Unlock()
	return nil, nil
}

// PartitionQuery is a placeholder for SpannerServer.PartitionQuery.
func (m *MockCloudSpannerClient) PartitionQuery(ctx context.Context, r *sppb.PartitionQueryRequest, opts ...grpc.CallOption) (*sppb.PartitionResponse, error) {
	m.ready()
	m.ReceivedRequests <- r

	return nil, errors.New("Unimplemented")
}

// PartitionRead is a placeholder for SpannerServer.PartitionRead.
func (m *MockCloudSpannerClient) PartitionRead(ctx context.Context, r *sppb.PartitionReadRequest, opts ...grpc.CallOption) (*sppb.PartitionResponse, error) {
	m.ready()
	m.ReceivedRequests <- r

	return nil, errors.New("Unimplemented")
}

// Freeze stalls all requests.
func (m *MockCloudSpannerClient) Freeze() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.freezed = make(chan struct{})
}

// Unfreeze restores processing requests.
func (m *MockCloudSpannerClient) Unfreeze() {
	m.mu.Lock()
	defer m.mu.Unlock()
	close(m.freezed)
}

// ready checks conditions before executing requests
// TODO: add checks for injected errors, actions
func (m *MockCloudSpannerClient) ready() {
	m.mu.Lock()
	freezed := m.freezed
	m.mu.Unlock()
	// check if client should be freezed
	<-freezed
}