From 5362184c80f069c96a9310143fb53fb749a7a8bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Caner=20=C3=87=C4=B1dam?= <canercidam01@gmail.com>
Date: Mon, 19 Jun 2023 18:12:38 +0300
Subject: [PATCH] improve bot manager update test case

---
 .../components/lifecycle/bot_manager_test.go  | 33 ++++++++++++-------
 .../components/lifecycle/lifecycle_test.go    |  5 ++-
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/services/components/lifecycle/bot_manager_test.go b/services/components/lifecycle/bot_manager_test.go
index f88474d6..9307daba 100644
--- 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
index 981f393a..88e0418b 100644
--- 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"