From 4d41b089f10d74ce30be540f6a5489ed6c726b76 Mon Sep 17 00:00:00 2001
From: omerlavanet <omer@lavanet.xyz>
Date: Wed, 6 Mar 2024 15:27:06 +0200
Subject: [PATCH] refactor code for reuse in unitest

---
 protocol/integration/protocol_test.go | 66 +++++++++++++++------------
 1 file changed, 36 insertions(+), 30 deletions(-)

diff --git a/protocol/integration/protocol_test.go b/protocol/integration/protocol_test.go
index ebe9869b31..435ae35834 100644
--- a/protocol/integration/protocol_test.go
+++ b/protocol/integration/protocol_test.go
@@ -65,42 +65,17 @@ func checkServerStatusWithTimeout(url string, totalTimeout time.Duration) bool {
 
 	return false
 }
-func TestConsumerProviderBasic(t *testing.T) {
-	ctx := context.Background()
+
+func createRpcConsumer(t *testing.T, ctx context.Context, specId string, apiInterface string, consumerListenAddress string, epoch uint64, pairingList map[uint64]*lavasession.ConsumerSessionsWithProvider, requiredResponses int, lavaChainID string) *rpcconsumer.RPCConsumerServer {
 	serverHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// Handle the incoming request and provide the desired response
 		w.WriteHeader(http.StatusOK)
 	})
-	// can be any spec and api interface
-	specId := "LAV1"
-	apiInterface := spectypes.APIInterfaceTendermintRPC
-	epoch := uint64(100)
-	requiredResponses := 1
-	lavaChainID := "lava"
-	providerListenAddress := "localhost:0"
-	consumerListenAddress := "localhost:11111"
-	pairingList := map[uint64]*lavasession.ConsumerSessionsWithProvider{
-		1: {
-			PublicLavaAddress: "",
-			Endpoints: []*lavasession.Endpoint{
-				{
-					NetworkAddress: providerListenAddress,
-					Enabled:        true,
-					Geolocation:    0,
-				},
-			},
-			Sessions:         map[int64]*lavasession.SingleConsumerSession{},
-			MaxComputeUnits:  10000,
-			UsedComputeUnits: 0,
-			PairingEpoch:     epoch,
-		},
-	}
-
 	chainParser, _, chainFetcher, _, err := chainlib.CreateChainLibMocks(ctx, specId, apiInterface, serverHandler, "../../", nil)
 	require.NoError(t, err)
 	require.NotNil(t, chainParser)
 	require.NotNil(t, chainFetcher)
-	// create basic consumer stuff
+
 	specIdLava := "LAV1"
 	chainParserLava, _, chainFetcherLava, _, err := chainlib.CreateChainLibMocks(ctx, specIdLava, spectypes.APIInterfaceRest, serverHandler, "../../", nil)
 	require.NoError(t, err)
@@ -118,7 +93,7 @@ func TestConsumerProviderBasic(t *testing.T) {
 	consumerStateTracker := &mockConsumerStateTracker{}
 	finalizationConsensus := lavaprotocol.NewFinalizationConsensus(rpcEndpoint.ChainID)
 	_, averageBlockTime, _, _ := chainParser.ChainBlockStats()
-	baseLatency := common.AverageWorldLatency / 2 // we want performance to be half our timeout or better
+	baseLatency := common.AverageWorldLatency / 2
 	optimizer := provideroptimizer.NewProviderOptimizer(provideroptimizer.STRATEGY_BALANCED, averageBlockTime, baseLatency, 2)
 	consumerSessionManager := lavasession.NewConsumerSessionManager(rpcEndpoint, optimizer, nil, nil)
 
@@ -131,11 +106,42 @@ func TestConsumerProviderBasic(t *testing.T) {
 	require.NoError(t, err)
 	err = rpcConsumerServer.ServeRPCRequests(ctx, rpcEndpoint, consumerStateTracker, chainParser, finalizationConsensus, consumerSessionManager, requiredResponses, account.SK, lavaChainID, nil, rpcsonumerLogs, account.Addr, consumerConsistency, nil, consumerCmdFlags, false, nil, nil)
 	require.NoError(t, err)
-
 	// wait for consumer server to be up
 	consumerUp := checkServerStatusWithTimeout("http://"+consumerListenAddress, time.Millisecond*50)
 	require.True(t, consumerUp)
 
+	return rpcConsumerServer
+}
+
+func TestConsumerProviderBasic(t *testing.T) {
+	ctx := context.Background()
+	// can be any spec and api interface
+	specId := "LAV1"
+	apiInterface := spectypes.APIInterfaceTendermintRPC
+	epoch := uint64(100)
+	requiredResponses := 1
+	lavaChainID := "lava"
+	providerListenAddress := "localhost:0"
+	consumerListenAddress := "localhost:11111"
+	pairingList := map[uint64]*lavasession.ConsumerSessionsWithProvider{
+		1: {
+			PublicLavaAddress: "",
+			Endpoints: []*lavasession.Endpoint{
+				{
+					NetworkAddress: providerListenAddress,
+					Enabled:        true,
+					Geolocation:    0,
+				},
+			},
+			Sessions:         map[int64]*lavasession.SingleConsumerSession{},
+			MaxComputeUnits:  10000,
+			UsedComputeUnits: 0,
+			PairingEpoch:     epoch,
+		},
+	}
+
+	rpcconsumerServer := createRpcConsumer(t, ctx, specId, apiInterface, consumerListenAddress, epoch, pairingList, requiredResponses, lavaChainID)
+	require.NotNil(t, rpcconsumerServer)
 	client := http.Client{}
 	resp, err := client.Get("http://" + consumerListenAddress + "/status")
 	require.NoError(t, err)