fix unit test

Author: YixuanWang-99

tests/unit/elastic_utils_test.py

@@ -341,6 +341,7 @@ def test_record_elastic_event_start(self):
     """Tests recording an elastic slice down start."""
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.new_slice_event.is_set.return_value = False
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
     config = FakeConfig()
 
@@ -356,6 +357,7 @@ def test_record_elastic_event_start_scale_up(self):
     """Tests recording an elastic slice scale up start."""
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.new_slice_event.is_set.return_value = True
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
     config = FakeConfig()
 
@@ -380,6 +382,9 @@ def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self):
   def test_record_elastic_wait_end_and_reinit_start(self):
     """Test recording end of slice down and start of reinit."""
     elastic_utils.pending_elastic_event_type = 'elastic_slice_down'
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_manager.active_slice_indices = {0}
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
 
     elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder)
@@ -398,6 +403,9 @@ def test_record_elastic_reinit_end(self):
     """Tests recording end of elastic reinit."""
     fake_recorder = Mock()
     elastic_utils.pending_reinit_recorder = fake_recorder
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_manager.active_slice_indices = {0}
+    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
 
     elastic_utils.record_elastic_reinit_end()