Selenium waits on message to be handled

Removes flakiness on fast machine

Author: vladimirnani

tests/sample_project/sampleapp/static/sampleapp/js/scripts.js

@@ -12,6 +12,7 @@
     const socket = new WebSocket(wsPath);
 
     window.websocketConnected = false;
+    window.messageHandled = false;
 
     socket.onopen = () => {
       window.websocketConnected = true;
@@ -27,6 +28,7 @@
   function handleMessage(e) {
     const data = JSON.parse(e.data);
     renderState(data.count, data.messages);
+    window.messageHandled = true;
   }
 
   function renderState(count, messages) {

tests/sample_project/tests/selenium_mixin.py

@@ -141,6 +141,21 @@ def wait_for_websocket_connection(self, substring="WebSocket connected", timeout
                 f"Timed out waiting for window.websocketConnected after {timeout}s"
             )
 
+    def wait_for_websocket_message_handled(self, timeout=5):
+        """
+        Wait until window.messageHandled is true, or fail after timeout.
+        """
+        try:
+            WebDriverWait(self.web_driver, timeout).until(
+                lambda d: d.execute_script("return window.messageHandled === true")
+            )
+        except TimeoutException:
+            logs = self.get_browser_logs()
+            print("\n Browser logs on WS-flag timeout:")
+            for entry in logs:
+                print(f"[{entry['level']}] {entry['message']}")
+            self.fail(f"Timed out waiting for window.messageHandled after {timeout}s")
+
     def tearDown(self):
         logs = self.web_driver.get_log("browser")
         severe_logs = [entry for entry in logs if entry.get("level") == "SEVERE"]

tests/sample_project/tests/test_selenium.py

@@ -13,6 +13,9 @@ class TestSampleApp(SeleniumMixin, ChannelsLiveServerTestCase):
     def setUp(self):
         super().setUp()
         self.login()
+        self.open_admin_message_page()
+
+    def open_admin_message_page(self):
         self.open("/admin/sampleapp/message/")
         self.wait_for_websocket_connection()
 
@@ -29,7 +32,7 @@ def test_real_time_create_message(self):
         tabs = self.web_driver.window_handles
         self.web_driver.switch_to.window(tabs[1])
 
-        self.open("/admin/sampleapp/message/")
+        self.open_admin_message_page()
         titleInput = self.find_element(By.ID, "msgTitle")
         self.assertIsNotNone(titleInput, "Title input should be present")
         messageInput = self.find_element(By.ID, "msgTextArea")
@@ -52,6 +55,7 @@ def test_real_time_create_message(self):
     def test_real_time_delete_message(self):
         self._create_message()
         self.web_driver.refresh()
+        self.wait_for_websocket_message_handled()
 
         messageCount = self.find_element(By.ID, "messageCount")
         self.assertIsNotNone(messageCount, "Message count should be present")
@@ -61,7 +65,7 @@ def test_real_time_delete_message(self):
         tabs = self.web_driver.window_handles
         self.web_driver.switch_to.window(tabs[1])
 
-        self.open("/admin/sampleapp/message/")
+        self.open_admin_message_page()
         deleteButton = self.find_element(By.ID, "deleteBtn")
         self.assertIsNotNone(deleteButton, "Delete button should be present")
         deleteButton.click()