Merge pull request #772 from forta-network/caner/improve-bot-manager-…

…update-test-case Improve bot manager update test case
forta-network · Jun 19, 2023 · 027dbff · 027dbff
2 parents 61a6bdc + 5362184
commit 027dbff
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 13 deletions.
diff --git a/services/components/lifecycle/bot_manager_test.go b/services/components/lifecycle/bot_manager_test.go
@@ -60,40 +60,49 @@ func (s *BotLifecycleManagerTestSuite) TestAddUpdateRemove() {
 	alreadyRunning := []config.AgentConfig{
 		{
 			ID:    testBotID1,
-			Image: testImageRef,
+			Image: testImageRef1,
+			ShardConfig: &config.ShardConfig{
+				ShardID: 0,
+				Shards:  2,
+				Target:  1,
+			},
 		},
 		{
 			ID:    testBotID2,
-			Image: testImageRef,
+			Image: testImageRef2,
 		},
 	}
 	latestAssigned := []config.AgentConfig{
 		{
 			ID:    testBotID3,
-			Image: testImageRef,
+			Image: testImageRef3,
 		},
 		{
 			ID:    testBotID1,
-			Image: testImageRef,
+			Image: testImageRef1,
 			ShardConfig: &config.ShardConfig{
-				ShardID: 1,
+				ShardID: 1, // shard config update
+				Shards:  2,
+				Target:  1,
 			},
 		},
 	}
-	addedBot := latestAssigned[0]
-	removedBot := alreadyRunning[1]
+	removedBots := alreadyRunning
+	addedBots := latestAssigned
 
 	s.botManager.runningBots = alreadyRunning
 
 	s.botRegistry.EXPECT().LoadAssignedBots().Return(latestAssigned, nil).Times(1)
 	s.lifecycleMetrics.EXPECT().SystemStatus("load.assigned.bots", "2")
 
-	s.botContainers.EXPECT().EnsureBotImages(gomock.Any(), []config.AgentConfig{addedBot}).Return([]error{nil}).Times(1)
-	s.botContainers.EXPECT().LaunchBot(gomock.Any(), addedBot).Return(nil).Times(1)
+	s.botPool.EXPECT().RemoveBotsWithConfigs(removedBots)
+	s.lifecycleMetrics.EXPECT().StatusStopping(removedBots)
+	s.botContainers.EXPECT().TearDownBot(gomock.Any(), removedBots[0].ContainerName(), true)
+	s.botContainers.EXPECT().TearDownBot(gomock.Any(), removedBots[1].ContainerName(), true)
 
-	s.botPool.EXPECT().RemoveBotsWithConfigs([]config.AgentConfig{removedBot})
-	s.lifecycleMetrics.EXPECT().StatusStopping([]config.AgentConfig{removedBot})
-	s.botContainers.EXPECT().TearDownBot(gomock.Any(), removedBot.ContainerName(), true)
+	s.botContainers.EXPECT().EnsureBotImages(gomock.Any(), addedBots).Return([]error{nil, nil}).Times(1)
+	s.botContainers.EXPECT().LaunchBot(gomock.Any(), addedBots[0]).Return(nil).Times(1)
+	s.botContainers.EXPECT().LaunchBot(gomock.Any(), addedBots[1]).Return(nil).Times(1)
 
 	s.lifecycleMetrics.EXPECT().StatusRunning(latestAssigned).Times(1)
 	s.botPool.EXPECT().UpdateBotsWithLatestConfigs(latestAssigned)

diff --git a/services/components/lifecycle/lifecycle_test.go b/services/components/lifecycle/lifecycle_test.go
@@ -26,7 +26,10 @@ const (
 	testBotID1       = "0x0100000000000000000000000000000000000000000000000000000000000000"
 	testBotID2       = "0x0200000000000000000000000000000000000000000000000000000000000000"
 	testBotID3       = "0x0300000000000000000000000000000000000000000000000000000000000000"
-	testImageRef     = "bafybeielvnt5apaxbk6chthc4dc3p6vscpx3ai4uvti7gwh253j7facsxu@sha256:e0e9efb6699b02750f6a9668084d37314f1de3a80da7e19c1d40da73ee57dd45"
+	testImageRef1    = "test-image-ref-1"
+	testImageRef2    = "test-image-ref-2"
+	testImageRef3    = "test-image-ref-3"
+	testImageRef     = testImageRef1
 	testContainerID  = "test-container-id"
 	testContainerID1 = "test-container-id-1"
 	testContainerID2 = "test-container-id-2"