Restore netplan config even when exception is raised (bugfix)

providers/base/bin/wifi_client_test_netplan.py

@@ -25,6 +25,8 @@
 import ipaddress
 import yaml
 
+from contextlib import contextmanager
+
 from gateway_ping_test import ping
 from checkbox_support.snap_utils.system import get_series
 
@@ -218,7 +220,10 @@
 
 def delete_test_config():
     print_head("Delete the test file")
-    os.remove(NETPLAN_TEST_CFG)
+    try:
+        os.remove(NETPLAN_TEST_CFG)
+    except FileNotFoundError:
+        print("Test config {} not found, ignoring...".format(NETPLAN_TEST_CFG))
     print()
 
 
@@ -232,11 +237,9 @@
     env.pop("PYTHONUSERBASE", None)
     retcode = sp.call(cmd, shell=True, env=env)
     if retcode != 0:
-        print("ERROR: failed netplan apply call")
-        print()
-        return False
+        print_journal_entries(start_time, renderer)
+        raise SystemExit("ERROR: failed netplan apply call")
     print()
-    return True
 
 
 def get_interface_info(interface, renderer):
@@ -450,65 +453,69 @@
     return parser.parse_args()
 
 
-def main():
-    args = parse_args()
-
-    start_time = datetime.datetime.now()
-
-    renderer = check_and_get_renderer(args.renderer)
-    args.renderer = renderer
+@contextmanager
+def handle_original_np_config():
     netplan_config_backup()
     netplan_config_wipe()
+    try:
+        yield
+    finally:
+        delete_test_config()
+        netplan_config_restore()
+        netplan_apply_config()
+
 
-    # Create wireless network test configuration file
+@contextmanager
+def handle_test_np_config(args):
     print_head("Generate a test netplan configuration")
     config_data = generate_test_config(**vars(args))
     print(config_data)
     print()
-
     write_test_config(config_data)
 
     # Bring up the interface
     print_head("Apply the test configuration")
-    if not netplan_apply_config():
-        delete_test_config()
-        netplan_config_restore()
-        print_journal_entries(start_time, renderer)
-        raise SystemExit(1)
+    netplan_apply_config()
     time.sleep(20)
+    try:
+        yield
+    finally:
+        delete_test_config()
 
-    print_head("Wait for interface to be routable")
-    reached_routable = wait_for_routable(args.interface, renderer)
-
-    test_result = False
-    if reached_routable:
-        print_head("Display address")
-        print_address_info(args.interface)
-
-        print_head("Display route table")
-        print_route_info()
-
-        # Check connection by ping or link status
-        print_head("Perform a ping test")
-        test_result = perform_ping_test(args.interface, renderer)
-        if test_result:
-            print("Connection test passed\n")
-        else:
-            print("Connection test failed\n")
 
-    delete_test_config()
-    netplan_config_restore()
+def main():
+    args = parse_args()
 
-    if not netplan_apply_config():
-        print_journal_entries(start_time, renderer)
-        raise SystemExit("ERROR: failed to apply restored config")
+    global start_time
+    global renderer
+    start_time = datetime.datetime.now()
+    renderer = check_and_get_renderer(args.renderer)
+    args.renderer = renderer
 
+    with handle_original_np_config():
+        with handle_test_np_config(args):
+            print_head("Wait for interface to be routable")
+            reached_routable = wait_for_routable(args.interface, renderer)
+
+            test_result = False
+            if reached_routable:
+                print_head("Display address")
+                print_address_info(args.interface)
+
+                print_head("Display route table")
+                print_route_info()
+
+                # Check connection by ping or link status
+                print_head("Perform a ping test")
+                test_result = perform_ping_test(args.interface, renderer)
+                if test_result:
+                    print("Connection test passed\n")
+                else:
+                    print("Connection test failed\n")
+    print_journal_entries(start_time, renderer)
     if not test_result:
-        print_journal_entries(start_time, renderer)
         raise SystemExit(1)
 
-    print_journal_entries(start_time, renderer)
-
 
 if __name__ == "__main__":
     # Check if executed with root

providers/base/tests/test_wifi_client_test_netplan.py

@@ -822,7 +822,6 @@ def test_main_success(
         mock_renderer.return_value = "networkd"
         mock_wait_routable.return_value = True
         mock_ping.return_value = True
-        mock_apply.return_value = True
 
         # Execute
         main()
@@ -839,7 +838,7 @@ def test_main_success(
         mock_print_address.assert_called_once_with("wlan0")
         self.assertEqual(mock_print_route.call_count, 1)
         mock_ping.assert_called_once_with("wlan0", "networkd")
-        self.assertEqual(mock_delete.call_count, 1)
+        self.assertEqual(mock_delete.call_count, 2)
         self.assertEqual(mock_restore.call_count, 1)
         self.assertEqual(mock_print_journal.call_count, 1)
 
@@ -849,20 +848,20 @@ def test_main_success(
     @patch("wifi_client_test_netplan.netplan_config_wipe")
     @patch("wifi_client_test_netplan.generate_test_config")
     @patch("wifi_client_test_netplan.write_test_config")
-    @patch("wifi_client_test_netplan.netplan_apply_config")
     @patch("wifi_client_test_netplan.time.sleep")
     @patch("wifi_client_test_netplan.wait_for_routable")
     @patch("wifi_client_test_netplan.delete_test_config")
     @patch("wifi_client_test_netplan.netplan_config_restore")
     @patch("wifi_client_test_netplan.print_journal_entries")
+    @patch("wifi_client_test_netplan.sp.call")
     def test_main_apply_config_failure(
         self,
+        mock_sp_call,
         mock_print_journal,
         mock_restore,
         mock_delete,
         mock_wait_routable,
         mock_sleep,
-        mock_apply,
         mock_write,
         mock_generate,
         mock_wipe,
@@ -875,15 +874,15 @@ def test_main_apply_config_failure(
         mock_args.renderer = "networkd"
         mock_parse_args.return_value = mock_args
         mock_renderer.return_value = "networkd"
-        mock_apply.return_value = False
+        mock_sp_call.return_value = 1  # Simulating an error
 
         # Execute and Assert
         with self.assertRaises(SystemExit):
             main()
 
         self.assertEqual(mock_delete.call_count, 1)
         self.assertEqual(mock_restore.call_count, 1)
-        self.assertEqual(mock_print_journal.call_count, 1)
+        self.assertEqual(mock_print_journal.call_count, 2)
 
     @patch("wifi_client_test_netplan.parse_args")
     @patch("wifi_client_test_netplan.check_and_get_renderer")
@@ -921,14 +920,14 @@ def test_main_ping_test_failure(
         mock_parse_args.return_value = mock_args
         mock_renderer.return_value = "networkd"
         mock_wait_routable.return_value = True
-        mock_apply.return_value = True
         mock_ping.return_value = False
 
         # Execute and Assert
         with self.assertRaises(SystemExit):
             main()
 
-        self.assertEqual(mock_delete.call_count, 1)
+        self.assertEqual(mock_delete.call_count, 2)
         self.assertEqual(mock_restore.call_count, 1)
+        self.assertEqual(mock_apply.call_count, 2)
         self.assertEqual(mock_print_journal.call_count, 1)
         mock_ping.assert_called_once_with("wlan0", "networkd")