fix up tests

tier4 · Nov 21, 2024 · 99a4a47 · 99a4a47
1 parent 6b3c744
commit 99a4a47
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/tests/test_otaclient/test_boot_control/test_grub.py b/tests/test_otaclient/test_boot_control/test_grub.py
@@ -20,6 +20,7 @@
 import shutil
 import typing
 from pathlib import Path
+from typing import Any
 
 import pytest
 import pytest_mock
@@ -368,10 +369,12 @@ def test_grub_normal_update(
         shutil.copy(slot_a_ota_partition_dir / _initrd, slot_b / "boot")
 
         logger.info("pre-update completed, entering post-update...")
+
         # test post-update
-        _post_updater = grub_controller.post_update()
-        next(_post_updater)
-        next(_post_updater, None)
+        grub_controller: Any  # for typing
+        grub_controller.post_update()
+        grub_controller.finalizing_update()
+
         assert (
             slot_b / Path(cfg.FSTAB_FILE).relative_to("/")
         ).read_text().strip() == self.FSTAB_UPDATED.strip()

diff --git a/tests/test_otaclient/test_boot_control/test_rpi_boot.py b/tests/test_otaclient/test_boot_control/test_rpi_boot.py
@@ -19,6 +19,7 @@
 import typing
 from pathlib import Path
 from string import Template
+from typing import Any
 
 import pytest
 import pytest_mock
@@ -262,10 +263,10 @@ def test_rpi_boot_normal_update(self, mocker: pytest_mock.MockerFixture):
         mocker.patch(f"{MODULE}.cfg", _mock_otaclient_cfg)
 
         # ------ boot_controller_inst1.stage1: init ------ #
-        rpi_boot_controller1 = RPIBootController()
+        rpi_boot_controller = RPIBootController()
 
         # ------ boot_controller_inst1.stage2: pre_update ------ #
-        rpi_boot_controller1.pre_update(
+        rpi_boot_controller.pre_update(
             version=VERSION,
             standby_as_ref=False,
             erase_standby=False,
@@ -299,9 +300,9 @@ def test_rpi_boot_normal_update(self, mocker: pytest_mock.MockerFixture):
         shutil.copy(os.path.realpath(_initrd_img), self.slot_b_boot_dir)
 
         # ------ boot_controller_inst1.stage3: post_update, reboot switch boot ------ #
-        _post_updater = rpi_boot_controller1.post_update()
-        next(_post_updater)
-        next(_post_updater, None)  # actual reboot here
+        rpi_boot_controller: Any  # for typing only
+        rpi_boot_controller.post_update()
+        rpi_boot_controller.finalizing_update()
 
         # --- assertion --- #
         self.reboot_tryboot_mock.assert_called_once()