Merge branch 'main' into main

Author: DocGarbanzo

donkeycar/tests/test_web_socket.py

@@ -1,13 +1,10 @@
 from tornado import testing
 import tornado.websocket
 import tornado.web
-import tornado.ioloop
 import json
 from unittest.mock import Mock
 from donkeycar.parts.web_controller.web import WebSocketCalibrateAPI
-from time import sleep
-
-SLEEP = 0.5
+import tornado.gen
 
 
 class WebSocketCalibrateTest(testing.AsyncHTTPTestCase):
@@ -24,6 +21,17 @@ def get_app(self):
 
     def get_ws_url(self):
         return "ws://localhost:" + str(self.get_http_port()) + "/"
+    
+    async def wait_for_attribute_value(self, obj, attr_name, 
+                                       expected_value, timeout_seconds=5):
+        """Poll until an object's attribute equals the expected value or timeout."""
+        iterations = int(timeout_seconds / 0.1)
+        for _ in range(iterations):
+            if (hasattr(obj, attr_name) 
+                and getattr(obj, attr_name) == expected_value):
+                return True
+            await tornado.gen.sleep(0.1)
+        return False
 
     @tornado.testing.gen_test
     def test_calibrate_servo_esc_1b(self):
@@ -40,7 +48,10 @@ def test_calibrate_servo_esc_1b(self):
         data = {"config": {"STEERING_LEFT_PWM": 444}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train['steering'], 'left_pulse', 444)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train['steering'].left_pulse == 444
         assert self.app.drive_train['steering'].right_pulse is None
 
@@ -59,7 +70,10 @@ def test_calibrate_servo_esc_1a(self):
         data = {"config": {"STEERING_LEFT_PWM": 444}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train['steering'], 'left_pulse', 444)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train['steering'].left_pulse == 444
         assert self.app.drive_train['steering'].right_pulse is None
 
@@ -78,7 +92,10 @@ def test_calibrate_servo_esc_2b(self):
         data = {"config": {"STEERING_RIGHT_PWM": 555}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train['steering'], 'right_pulse', 555)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train['steering'].right_pulse == 555
         assert self.app.drive_train['steering'].left_pulse is None
 
@@ -97,7 +114,10 @@ def test_calibrate_servo_esc_2a(self):
         data = {"config": {"STEERING_RIGHT_PWM": 555}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train['steering'], 'right_pulse', 555)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train['steering'].right_pulse == 555
         assert self.app.drive_train['steering'].left_pulse is None
 
@@ -116,7 +136,10 @@ def test_calibrate_servo_esc_3b(self):
         data = {"config": {"THROTTLE_FORWARD_PWM": 666}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train['throttle'], 'max_pulse', 666)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train['throttle'].max_pulse == 666
         assert self.app.drive_train['throttle'].min_pulse is None
 
@@ -135,7 +158,10 @@ def test_calibrate_servo_esc_3a(self):
         data = {"config": {"THROTTLE_FORWARD_PWM": 666}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train['throttle'], 'max_pulse', 666)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train['throttle'].max_pulse == 666
         assert self.app.drive_train['throttle'].min_pulse is None
 
@@ -151,5 +177,8 @@ def test_calibrate_mm1(self):
         data = {"config": {"MM1_STEERING_MID": 1234}}
         yield ws_client.write_message(json.dumps(data))
         yield ws_client.close()
-        sleep(SLEEP)
+        
+        result = yield self.wait_for_attribute_value(
+            self.app.drive_train, 'STEERING_MID', 1234)
+        assert result, "WebSocket message not processed within timeout"
         assert self.app.drive_train.STEERING_MID == 1234