From 518350d24bca734c3e41e327c5a8e33661b3ff6d Mon Sep 17 00:00:00 2001
From: Bolek <1416262+bolekk@users.noreply.github.com>
Date: Tue, 23 Jul 2024 05:25:30 -0700
Subject: [PATCH] Ukflake TestFunctionsConnectorHandler (#13916)

Wait for the handler to be called before proceeding.
---
 core/services/functions/connector_handler_test.go | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/core/services/functions/connector_handler_test.go b/core/services/functions/connector_handler_test.go
index 7e1f3ced34f..a24139da58e 100644
--- a/core/services/functions/connector_handler_test.go
+++ b/core/services/functions/connector_handler_test.go
@@ -269,7 +269,10 @@ func TestFunctionsConnectorHandler(t *testing.T) {
 		// first call to trigger the request
 		var response functions.HeartbeatResponse
 		allowlist.On("Allow", addr).Return(true).Once()
-		listener.On("HandleOffchainRequest", mock.Anything, mock.Anything).Return(nil).Once()
+		handlerCalled := make(chan struct{})
+		listener.On("HandleOffchainRequest", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
+			handlerCalled <- struct{}{}
+		}).Return(nil).Once()
 		connector.On("SendToGateway", mock.Anything, "gw1", mock.Anything).Run(func(args mock.Arguments) {
 			respMsg, ok := args[2].(*api.Message)
 			require.True(t, ok)
@@ -277,6 +280,7 @@ func TestFunctionsConnectorHandler(t *testing.T) {
 			require.Equal(t, functions.RequestStatePending, response.Status)
 		}).Return(nil).Once()
 		handler.HandleGatewayMessage(ctx, "gw1", msg)
+		<-handlerCalled
 
 		// async response computation
 		reportCh <- &functions.OffchainResponse{
@@ -304,9 +308,13 @@ func TestFunctionsConnectorHandler(t *testing.T) {
 		// first call to trigger the request
 		var response functions.HeartbeatResponse
 		allowlist.On("Allow", addr).Return(true).Once()
-		listener.On("HandleOffchainRequest", mock.Anything, mock.Anything).Return(errors.New("boom")).Once()
+		handlerCalled := make(chan struct{})
+		listener.On("HandleOffchainRequest", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
+			handlerCalled <- struct{}{}
+		}).Return(errors.New("boom")).Once()
 		connector.On("SendToGateway", mock.Anything, "gw1", mock.Anything).Return(nil).Once()
 		handler.HandleGatewayMessage(ctx, "gw1", msg)
+		<-handlerCalled
 
 		// collect the response - should eventually result in an internal error
 		gomega.NewGomegaWithT(t).Eventually(func() bool {