From dc4734253a82d3273a39c304759edcc19979c81e Mon Sep 17 00:00:00 2001
From: isaacyang <isaac.yang@canonical.com>
Date: Mon, 16 Dec 2024 10:18:16 +0800
Subject: [PATCH] Add unit test

---
 tests/generator/test_wifis.py | 54 +++++++++++++++++++++++++++++++++++
 tests/parser/test_keyfile.py  | 50 ++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/tests/generator/test_wifis.py b/tests/generator/test_wifis.py
index 5809b40e7..a351fa387 100644
--- a/tests/generator/test_wifis.py
+++ b/tests/generator/test_wifis.py
@@ -275,6 +275,27 @@ def test_wifi_wowlan_default(self):
         self.assertTrue(os.path.islink(os.path.join(
             self.workdir.name, 'run/systemd/system/systemd-networkd.service.wants/netplan-wpa-wl0.service')))
 
+    def test_wifi_wpa_sha256(self):
+        self.generate('''network:
+  version: 2
+  wifis:
+    wl0:
+      access-points:
+        homenet:
+          auth:
+            key-management: psk-sha256
+            password: "********"''')
+
+        self.assert_wpa_supplicant("wl0", """ctrl_interface=/run/wpa_supplicant
+
+network={
+  ssid=P"homenet"
+  key_mgmt=WPA-PSK WPA-PSK-SHA256
+  ieee80211w=1
+  psk="********"
+}
+""")
+
     def test_wifi_wpa3_personal(self):
         self.generate('''network:
   version: 2
@@ -794,6 +815,39 @@ def test_wifi_adhoc_wpa_5ghz(self):
 }
 """)
 
+    def test_wifi_wpa2_sha256(self):
+        self.generate('''network:
+  version: 2
+  renderer: NetworkManager
+  wifis:
+    wl0:
+      access-points:
+        homenet:
+          auth:
+            key-management: psk-sha256
+            password: "********"''')
+
+        self.assert_nm({'wl0-homenet': '''[connection]
+id=netplan-wl0-homenet
+type=wifi
+interface-name=wl0
+
+[ipv4]
+method=link-local
+
+[ipv6]
+method=ignore
+
+[wifi]
+ssid=homenet
+mode=infrastructure
+
+[wifi-security]
+key-mgmt=wpa-psk
+pmf=1
+psk=********
+'''})
+
     def test_wifi_wpa3_personal(self):
         self.generate('''network:
   version: 2
diff --git a/tests/parser/test_keyfile.py b/tests/parser/test_keyfile.py
index e178f7729..50ae9a9b6 100644
--- a/tests/parser/test_keyfile.py
+++ b/tests/parser/test_keyfile.py
@@ -1505,6 +1505,56 @@ def test_keyfile_wpa3_sae(self):
         name: "test2"
 '''.format(UUID, UUID)})
 
+    def test_keyfile_wpa_sha256(self):
+        self.generate_from_keyfile('''[connection]
+id=test2
+uuid={}
+type=wifi
+interface-name=wlan0
+
+[wifi]
+mode=infrastructure
+ssid=ubuntu-wpa-sha256
+
+[wifi-security]
+key-mgmt=wpa-psk
+psk=test1234
+
+[ipv4]
+method=auto
+
+[ipv6]
+addr-gen-mode=stable-privacy
+method=auto
+
+[proxy]
+'''.format(UUID))
+        self.assert_netplan({UUID: '''network:
+  version: 2
+  wifis:
+    NM-{}:
+      renderer: NetworkManager
+      match:
+        name: "wlan0"
+      dhcp4: true
+      dhcp6: true
+      ipv6-address-generation: "stable-privacy"
+      access-points:
+        "ubuntu-wpa-sha256":
+          auth:
+            key-management: "wpa-psk-sha256"
+            password: "test1234"
+          networkmanager:
+            uuid: "ff9d6ebc-226d-4f82-a485-b7ff83b9607f"
+            name: "test2"
+            passthrough:
+              ipv6.ip6-privacy: "-1"
+              proxy._: ""
+      networkmanager:
+        uuid: "{}"
+        name: "test2"
+'''.format(UUID, UUID)})
+
     def test_keyfile_wpa3_enterprise_eap_sha256(self):
         self.generate_from_keyfile('''[connection]
 id=test2