esp_prov: Refactor to use new 'wait_wifi_connected' function

Means all provisioning examples will have the same retry behaviour.

Author: projectgus

examples/provisioning/legacy/ble_prov/ble_prov_test.py

@@ -17,7 +17,6 @@
 from __future__ import print_function
 import re
 import os
-import time
 
 import ttfw_idf
 import esp_prov
@@ -82,19 +81,7 @@ def test_examples_provisioning_ble(env, extra_data):
     if not esp_prov.apply_wifi_config(transport, security):
         raise RuntimeError("Failed to send apply config")
 
-    success = False
-    while True:
-        time.sleep(5)
-        print("Wi-Fi connection state")
-        ret = esp_prov.get_wifi_config(transport, security)
-        if (ret == "connecting"):
-            continue
-        elif (ret == "connected"):
-            print("Provisioning was successful")
-            success = True
-        break
-
-    if not success:
+    if not esp_prov.wait_wifi_connected(transport, security):
         raise RuntimeError("Provisioning failed")
 
 

examples/provisioning/legacy/softap_prov/softap_prov_test.py

@@ -17,7 +17,6 @@
 from __future__ import print_function
 import re
 import os
-import time
 
 import ttfw_idf
 import esp_prov
@@ -96,19 +95,7 @@ def test_examples_provisioning_softap(env, extra_data):
     if not esp_prov.apply_wifi_config(transport, security):
         raise RuntimeError("Failed to send apply config")
 
-    success = False
-    while True:
-        time.sleep(5)
-        print("Wi-Fi connection state")
-        ret = esp_prov.get_wifi_config(transport, security)
-        if (ret == "connecting"):
-            continue
-        elif (ret == "connected"):
-            print("Provisioning was successful")
-            success = True
-        break
-
-    if not success:
+    if not esp_prov.wait_wifi_connected(transport, security):
         raise RuntimeError("Provisioning failed")
 
 

examples/provisioning/wifi_prov_mgr/wifi_prov_mgr_test.py

@@ -17,7 +17,6 @@
 from __future__ import print_function
 import re
 import os
-import time
 
 import ttfw_idf
 import esp_prov
@@ -90,24 +89,7 @@ def test_examples_wifi_prov_mgr(env, extra_data):
     if not esp_prov.apply_wifi_config(transport, security):
         raise RuntimeError("Failed to send apply config")
 
-    success = False
-    retry = 0
-    while True:
-        time.sleep(5)
-        print("Wi-Fi connection state")
-        ret = esp_prov.get_wifi_config(transport, security)
-        if (ret == "connecting"):
-            continue
-        elif (ret == "connected"):
-            print("Provisioning was successful")
-            success = True
-        elif (ret == "failed" and retry < 3):
-            retry = retry + 1
-            print("Connection failed.. retry again...: ", ret)
-            continue
-        break
-
-    if not success:
+    if not esp_prov.wait_wifi_connected(transport, security):
         raise RuntimeError("Provisioning failed")
 
     # Check if BTDM memory is released after provisioning finishes

tools/esp_prov/esp_prov.py

@@ -276,6 +276,32 @@ def get_wifi_config(tp, sec):
         return None
 
 
+def wait_wifi_connected(tp, sec):
+    """
+    Wait for provisioning to report Wi-Fi is connected
+
+    Returns True if Wi-Fi connection succeeded, False if connection consistently failed
+    """
+    TIME_PER_POLL = 5
+    retry = 3
+
+    while True:
+        time.sleep(TIME_PER_POLL)
+        print("\n==== Wi-Fi connection state  ====")
+        ret = get_wifi_config(tp, sec)
+        if ret == "connecting":
+            continue
+        elif ret == "connected":
+            print("==== Provisioning was successful ====")
+            return True
+        elif retry > 0:
+            retry -= 1
+            print("Waiting to poll status again (status %s, %d tries left)..." % (ret, retry))
+        else:
+            print("---- Provisioning failed ----")
+            return False
+
+
 def desc_format(*args):
     desc = ''
     for arg in args:
@@ -471,14 +497,4 @@ def desc_format(*args):
         exit(7)
     print("==== Apply config sent successfully ====")
 
-    while True:
-        time.sleep(5)
-        print("\n==== Wi-Fi connection state  ====")
-        ret = get_wifi_config(obj_transport, obj_security)
-        if (ret == "connecting"):
-            continue
-        elif (ret == "connected"):
-            print("==== Provisioning was successful ====")
-        else:
-            print("---- Provisioning failed ----")
-        break
+    wait_wifi_connected(obj_transport, obj_security)

tools/esp_prov/prov/wifi_prov.py

@@ -65,7 +65,6 @@ def config_get_status_response(security_ctx, response_data):
     return "unknown"
 
 
-
 def config_set_config_request(security_ctx, ssid, passphrase):
     # Form protobuf request packet for SetConfig command
     cmd = proto.wifi_config_pb2.WiFiConfigPayload()