fix

pg-sharding · Apr 2, 2024 · a5f9629 · a5f9629
1 parent 0c36cf0
commit a5f9629
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/router/frontend/frontend_test.go b/router/frontend/frontend_test.go
@@ -88,7 +88,7 @@ func TestFrontendSimple(t *testing.T) {
 	cl.EXPECT().RUnlock().AnyTimes()
 
 	// reroute on first query in this case
-	cmngr.EXPECT().ValidateReRoute(gomock.Any()).Return(true)
+	cmngr.EXPECT().ValidateReRoute(gomock.Any()).AnyTimes().Return(true)
 
 	cmngr.EXPECT().RouteCB(cl, gomock.Any()).AnyTimes()
 
@@ -210,7 +210,7 @@ func TestFrontendXProto(t *testing.T) {
 	cl.EXPECT().RUnlock().AnyTimes()
 
 	// reroute on first query in this case
-	cmngr.EXPECT().ValidateReRoute(gomock.Any()).Return(true)
+	cmngr.EXPECT().ValidateReRoute(gomock.Any()).AnyTimes().Return(true)
 
 	cmngr.EXPECT().RouteCB(cl, gomock.Any()).AnyTimes()
 
@@ -243,13 +243,17 @@ func TestFrontendXProto(t *testing.T) {
 	cl.EXPECT().Receive().Times(1).Return(&pgproto3.Sync{}, nil)
 
 	cl.EXPECT().StorePreparedStatement("stmtcache_1", "select 'Hello, world!'").Times(1).Return()
-	cl.EXPECT().PreparedStatementQueryByName("stmtcache_1").Return("select 'Hello, world!'")
+	cl.EXPECT().PreparedStatementQueryByName("stmtcache_1").AnyTimes().Return("select 'Hello, world!'")
 
 	cl.EXPECT().ServerAcquireUse().AnyTimes()
 	cl.EXPECT().ServerReleaseUse().AnyTimes()
 
-	srv.EXPECT().HasPrepareStatement(gomock.Any()).Return(false, shard.PreparedStatementDescriptor{}).AnyTimes()
-	srv.EXPECT().PrepareStatement(gomock.Any(), gomock.Any()).AnyTimes()
+	res := false
+
+	srv.EXPECT().HasPrepareStatement(gomock.Any()).DoAndReturn(func(interface{}) (interface{}, interface{}) { return res, shard.PreparedStatementDescriptor{} }).AnyTimes()
+	srv.EXPECT().PrepareStatement(gomock.Any(), gomock.Any()).Do(func(interface{}, interface{}) {
+		res = true
+	}).AnyTimes()
 	/* */
 
 	srv.EXPECT().Send(&pgproto3.Parse{
@@ -289,7 +293,7 @@ func TestFrontendXProto(t *testing.T) {
 	}, nil)
 
 	// receive this 4 msgs
-	cl.EXPECT().Send(gomock.Any()).Times(3).Return(nil)
+	cl.EXPECT().Send(gomock.Any()).Times(4).Return(nil)
 
 	cl.EXPECT().Receive().Times(1).Return(nil, io.EOF)
 
@@ -339,7 +343,7 @@ func TestFrontendSimpleCopyIn(t *testing.T) {
 	cl.EXPECT().RUnlock().AnyTimes()
 
 	// reroute on first query in this case
-	cmngr.EXPECT().ValidateReRoute(gomock.Any()).Return(true)
+	cmngr.EXPECT().ValidateReRoute(gomock.Any()).AnyTimes().Return(true)
 
 	cmngr.EXPECT().RouteCB(cl, gomock.Any()).AnyTimes()
 

diff --git a/router/relay/relay.go b/router/relay/relay.go
@@ -221,15 +221,15 @@ func (rst *RelayStateImpl) PrepareStatement(hash uint64, d server.PrepStmtDesc)
 
 	var retMsg pgproto3.BackendMessage
 
-	deplotyed := false
+	deployed := false
 
 	for _, msg := range unreplied {
 		spqrlog.Zero.Debug().Uint("client", rst.Client().ID()).Interface("type", msg).Msg("unreplied msg in prepare")
 		switch q := msg.(type) {
 		case *pgproto3.ParseComplete:
 			// skip
 			retMsg = msg
-			deplotyed = true
+			deployed = true
 		case *pgproto3.ErrorResponse:
 			retMsg = msg
 		case *pgproto3.NoData:
@@ -244,7 +244,7 @@ func (rst *RelayStateImpl) PrepareStatement(hash uint64, d server.PrepStmtDesc)
 		}
 	}
 
-	if deplotyed {
+	if deployed {
 		// dont need to complete relay because tx state didt changed
 		rst.Cl.Server().PrepareStatement(hash, rd)
 	}