[PATCH v2 4/6] tests: refactor Hostapd event checks

[Johannes Berg]

From: Johannes Berg <johannes.berg at intel.com>

Quite a bit of code is duplicated here, refactor it.
Again, part of the motiviation is async code later.

Signed-off-by: Johannes Berg <johannes.berg at intel.com>
---
 tests/hwsim/hostapd.py | 43 ++++++++++++++++++++----------------------
 1 file changed, 20 insertions(+), 23 deletions(-)

diff --git a/tests/hwsim/hostapd.py b/tests/hwsim/hostapd.py
index a0a2340fdf36..6ce6701667b9 100644
--- a/tests/hwsim/hostapd.py
+++ b/tests/hwsim/hostapd.py
@@ -293,37 +293,34 @@ class Hostapd:
         return wpaspy.wait_event(self.mon, events, timeout,
                                  log_prefix=self.dbg + ": ")
 
+    @staticmethod
+    def _check_sta_event(ev, addr, what):
+        if ev is None:
+            raise Exception("AP did not report " + what)
+        if addr and addr not in ev:
+            raise Exception("Unexpected STA address in " + what +
+                            " event: " + ev)
+        return ev
 
     def wait_sta(self, addr=None, timeout=2, wait_4way_hs=False):
-        ev = self.wait_event(["AP-STA-CONNECT"], timeout=timeout)
-        if ev is None:
-            raise Exception("AP did not report STA connection")
-        if addr and addr not in ev:
-            raise Exception("Unexpected STA address in connection event: " + ev)
+        ev = self._check_sta_event(
+            self.wait_event(["AP-STA-CONNECT"], timeout=timeout), addr,
+            "STA connection")
         if wait_4way_hs:
-            ev2 = self.wait_event(["EAPOL-4WAY-HS-COMPLETED"],
-                                  timeout=timeout)
-            if ev2 is None:
-                raise Exception("AP did not report 4-way handshake completion")
-            if addr and addr not in ev2:
-                raise Exception("Unexpected STA address in 4-way handshake completion event: " + ev2)
+            self._check_sta_event(
+                self.wait_event(["EAPOL-4WAY-HS-COMPLETED"], timeout=timeout),
+                addr, "4-way handshake completion")
         return ev
 
     def wait_sta_disconnect(self, addr=None, timeout=2):
-        ev = self.wait_event(["AP-STA-DISCONNECT"], timeout=timeout)
-        if ev is None:
-            raise Exception("AP did not report STA disconnection")
-        if addr and addr not in ev:
-            raise Exception("Unexpected STA address in disconnection event: " + ev)
-        return ev
+        return self._check_sta_event(
+            self.wait_event(["AP-STA-DISCONNECT"], timeout=timeout), addr,
+            "STA disconnection")
 
     def wait_4way_hs(self, addr=None, timeout=1):
-        ev = self.wait_event(["EAPOL-4WAY-HS-COMPLETED"], timeout=timeout)
-        if ev is None:
-            raise Exception("hostapd did not report 4-way handshake completion")
-        if addr and addr not in ev:
-            raise Exception("Unexpected STA address in 4-way handshake completion event: " + ev)
-        return ev
+        return self._check_sta_event(
+            self.wait_event(["EAPOL-4WAY-HS-COMPLETED"], timeout=timeout),
+            addr, "4-way handshake completion")
 
     def wait_ptkinitdone(self, addr, timeout=2):
         while timeout > 0:
-- 
2.55.0