From 89400705c66376bbe8a9eecbd1b9be392f10253a Mon Sep 17 00:00:00 2001
From: James Falcon <james.falcon@canonical.com>
Date: Thu, 17 Oct 2024 08:00:17 -0500
Subject: [PATCH] integration tests source uri check

---
 cloudinit/cmd/main.py                               | 3 ++-
 tests/integration_tests/datasources/test_nocloud.py | 4 ++--
 tests/integration_tests/modules/test_boothook.py    | 4 ++--
 tests/integration_tests/modules/test_combined.py    | 4 ++--
 tests/integration_tests/util.py                     | 7 +++++--
 5 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/cloudinit/cmd/main.py b/cloudinit/cmd/main.py
index 94efbf6ffee..a2851a2682c 100644
--- a/cloudinit/cmd/main.py
+++ b/cloudinit/cmd/main.py
@@ -362,7 +362,8 @@ def _should_wait_via_user_data(
     # These all have the potential to require network access, so we should wait
     if "write_files" in parsed_yaml:
         for item in parsed_yaml["write_files"]:
-            source_uri = item.get("source", {}).get("uri", "")
+            source_dict = item.get("source") or {}
+            source_uri = source_dict.get("uri", "")
             if source_uri and not (source_uri.startswith(("/", "file:"))):
                 return True, "write_files with source uri found"
         return False, "write_files without source uri found"
diff --git a/tests/integration_tests/datasources/test_nocloud.py b/tests/integration_tests/datasources/test_nocloud.py
index 0929e77914b..7380c37f00f 100644
--- a/tests/integration_tests/datasources/test_nocloud.py
+++ b/tests/integration_tests/datasources/test_nocloud.py
@@ -15,7 +15,7 @@
     override_kernel_command_line,
     verify_clean_boot,
     verify_clean_log,
-    wait_online_called,
+    network_wait_logged,
 )
 
 VENDOR_DATA = """\
@@ -100,7 +100,7 @@ def test_nocloud_seedfrom_vendordata(client: IntegrationInstance):
     client.restart()
     assert client.execute("cloud-init status").ok
     assert "seeded_vendordata_test_file" in client.execute("ls /var/tmp")
-    assert wait_online_called(client.execute("cat /var/log/cloud-init.log"))
+    assert network_wait_logged(client.execute("cat /var/log/cloud-init.log"))
 
 
 SMBIOS_USERDATA = """\
diff --git a/tests/integration_tests/modules/test_boothook.py b/tests/integration_tests/modules/test_boothook.py
index 4dbe266e2d2..55ca1ca6861 100644
--- a/tests/integration_tests/modules/test_boothook.py
+++ b/tests/integration_tests/modules/test_boothook.py
@@ -5,9 +5,9 @@
 
 from tests.integration_tests.instances import IntegrationInstance
 from tests.integration_tests.util import (
+    network_wait_logged,
     verify_clean_boot,
     verify_clean_log,
-    wait_online_called,
 )
 
 USER_DATA = """\
@@ -55,6 +55,6 @@ def test_boothook_waits_for_network(
     ):
         """Test boothook handling waits for network before running."""
         client = class_client
-        assert wait_online_called(
+        assert network_wait_logged(
             client.read_from_file("/var/log/cloud-init.log")
         )
diff --git a/tests/integration_tests/modules/test_combined.py b/tests/integration_tests/modules/test_combined.py
index 6a20ce3de04..9c525840eb6 100644
--- a/tests/integration_tests/modules/test_combined.py
+++ b/tests/integration_tests/modules/test_combined.py
@@ -30,9 +30,9 @@
     get_feature_flag_value,
     get_inactive_modules,
     lxd_has_nocloud,
+    network_wait_logged,
     verify_clean_boot,
     verify_ordered_items_in_text,
-    wait_online_called,
 )
 
 USER_DATA = """\
@@ -558,7 +558,7 @@ def test_unicode(self, class_client: IntegrationInstance):
     @pytest.mark.skipif(not IS_UBUNTU, reason="Ubuntu-only behavior")
     def test_networkd_wait_online(self, class_client: IntegrationInstance):
         client = class_client
-        assert not wait_online_called(
+        assert not network_wait_logged(
             client.read_from_file("/var/log/cloud-init.log")
         )
 
diff --git a/tests/integration_tests/util.py b/tests/integration_tests/util.py
index 7f561ae2325..f5f51dc48fa 100644
--- a/tests/integration_tests/util.py
+++ b/tests/integration_tests/util.py
@@ -634,5 +634,8 @@ def push_and_enable_systemd_unit(
     client.execute(f"systemctl enable {unit_name}", use_sudo=True)
 
 
-def wait_online_called(log: str) -> bool:
-    return "Running command ['/lib/systemd/systemd-networkd-wait-online" in log
+def network_wait_logged(log: str) -> bool:
+    return (
+        "Running command "
+        "['systemctl', 'start', 'systemd-networkd-wait-online.service']"
+    ) in log