From c3df3a62bde94aa76fe688b9433e9ca153a63ba3 Mon Sep 17 00:00:00 2001
From: Logan Lipke <lolipke@gmail.com>
Date: Mon, 22 Jun 2020 14:16:36 -0700
Subject: [PATCH] Improved cleanup function

---
 py-scripts/test_ipv4_connection.py | 56 ++++++++++++++++++------------
 1 file changed, 33 insertions(+), 23 deletions(-)

diff --git a/py-scripts/test_ipv4_connection.py b/py-scripts/test_ipv4_connection.py
index 65c4f0cb..e370b81d 100755
--- a/py-scripts/test_ipv4_connection.py
+++ b/py-scripts/test_ipv4_connection.py
@@ -10,16 +10,15 @@ if 'py-json' not in sys.path:
     sys.path.append('../py-json')
 
 import argparse
-from LANforge import LFUtils
-# from LANforge import LFCliBase
-from LANforge import lfcli_base
 from LANforge.lfcli_base import LFCliBase
 from LANforge.LFUtils import *
 import realm
 import time
 
+
 class IPv4Test(LFCliBase):
-    def __init__(self, host, port, ssid, security, password, num_stations, _debug_on=False, _exit_on_error=False,
+    def __init__(self, host, port, ssid, security, password, num_stations, prefix="00000", _debug_on=False,
+                 _exit_on_error=False,
                  _exit_on_fail=False):
         super().__init__(host, port, _debug=_debug_on, _halt_on_error=_exit_on_error, _exit_on_fail=_exit_on_fail)
         self.host = host
@@ -29,19 +28,31 @@ class IPv4Test(LFCliBase):
         self.password = password
         self.num_stations = num_stations
         self.timeout = 120
+        self.prefix = prefix
+        self.local_realm = realm.Realm(lfclient_host=self.host, lfclient_port=self.port)
+        self.profile = realm.StationProfile(self.lfclient_url, ssid=self.ssid, ssid_pass=self.password,
+                                            security=self.security, prefix=self.prefix, mode=0, up=True, dhcp=True,
+                                            debug_=False)
 
     def run_test(self, sta_list, print_pass=False, print_fail=False):
         for sta_name in sta_list:
             sta_status = super().json_get("port/1/1/" + sta_name)
-            if len(sta_status['interface']['ap']) != 18 and sta_status['interface']['ap'][-3] != ':':
+            if sta_status['interface']['ap'] == '' or len(sta_status['interface']['ap']) != 18 and \
+                    sta_status['interface']['ap'][-3] != ':':
                 self._fail("%s failed to associate" % sta_name, print_fail)
             elif sta_status['interface']['ip'] == "0.0.0.0":
                 self._fail("%s failed to get ip" % sta_name, print_fail)
             else:
                 self._pass("%s associated with ip" % sta_name, print_pass)
 
-    def cleanup(self, sta_list):
-        print("Cleaning up %s " % sta_list)
+    def cleanup(self):
+        port_list = self.local_realm.station_list()
+        sta_list = []
+        for item in list(port_list):
+            # print(list(item))
+            if "sta" in list(item)[0]:
+                sta_list.append(self.local_realm.name_to_eid(list(item)[0])[2])
+
         for sta_name in sta_list:
             req_url = "cli-json/rm_vlan"
             data = {
@@ -54,23 +65,21 @@ class IPv4Test(LFCliBase):
 
     def run(self):
         super().clear_test_results()
-        local_realm = realm.Realm(lfclient_host=self.host, lfclient_port=self.port)
-        profile = realm.StationProfile(self.lfclient_url, ssid=self.ssid, ssid_pass=self.password,
-                                             security=self.security, prefix="00000", mode=0, up=True, dhcp=True,
-                                             debug_=False)
-        profile.use_wpa2(True, self.ssid, self.password)
-        profile.set_prefix("000")
-        profile.create(resource = 1, radio = "wiphy0", num_stations = self.num_stations)
-        port_list = local_realm.station_list()
+        print("Cleaning up old stations")
+        self.cleanup()
         sta_list = []
-        for item in list(port_list):
-            if "sta" in item:
-                sta_list.append(local_realm.name_to_eid(item)[2])
 
-        for sec in range(self.timeout):
-            self.run_test(sta_list, True, True)
-            time.sleep(1)
-        self.cleanup(sta_list)
+        self.profile.use_wpa2(True, self.ssid, self.password)
+        self.profile.set_prefix(self.prefix)
+        self.profile.create(resource=1, radio="wiphy0", num_stations=self.num_stations)
+        for sta_name in list(self.local_realm.find_ports_like("sta[%s..%s]" % (self.prefix[:-1], str(self.prefix[:-2]) + str(self.num_stations - 1)))):
+            sta_list.append(self.local_realm.name_to_eid(sta_name)[2])
+        print(sta_list)
+        time.sleep(self.timeout)
+        self.run_test(sta_list, True, True)
+        print("Cleaning up")
+        self.cleanup()
+
 
 def main():
     lfjson_host = "localhost"
@@ -80,5 +89,6 @@ def main():
     ip_test.timeout = 120
     ip_test.run()
 
+
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()