Misc fixes and test robustness changes

Author: RaphaelRobidas

frontend/calcusliveserver.py

@@ -36,6 +36,7 @@
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.chrome.options import Options
 from selenium.webdriver.common.action_chains import ActionChains
+from selenium.webdriver.remote.client_config import ClientConfig
 
 from celery.contrib.testing.worker import start_worker
 from celery.contrib.abortable import AbortableAsyncResult
@@ -77,6 +78,7 @@
 
 class CalcusLiveServer(StaticLiveServerTestCase):
     host = "0.0.0.0"
+    REMOTE_DRIVER_TIMEOUT = 60
 
     @classmethod
     def _make_driver(cls):
@@ -85,9 +87,13 @@ def _make_driver(cls):
             chrome_options.add_argument("--headless")
             driver = webdriver.Chrome(options=chrome_options)
         else:
+            selenium_hub = "http://selenium:4444/wd/hub"
             driver = webdriver.Remote(
-                command_executor="http://selenium:4444/wd/hub",
+                command_executor=selenium_hub,
                 options=chrome_options,
+                client_config=ClientConfig(
+                    remote_server_addr=selenium_hub, timeout=cls.REMOTE_DRIVER_TIMEOUT
+                ),
             )
 
         driver.set_window_size(ZOOM * 1920, ZOOM * 1080)
@@ -102,6 +108,18 @@ def _restart_driver(cls):
             pass
         cls.driver = cls._make_driver()
 
+    @classmethod
+    def _restart_celery_worker(cls):
+        if IS_CLOUD:
+            return
+
+        if hasattr(cls, "celery_worker"):
+            cls.celery_worker.__exit__(None, None, None)
+
+        app.loader.import_module("celery.contrib.testing.tasks")
+        cls.celery_worker = start_worker(app, perform_ping_check=False)
+        cls.celery_worker.__enter__()
+
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
@@ -112,10 +130,7 @@ def setUpClass(cls):
         tasks.REMOTE = False
 
         if not IS_CLOUD:
-            app.loader.import_module("celery.contrib.testing.tasks")
-
-            cls.celery_worker = start_worker(app, perform_ping_check=False)
-            cls.celery_worker.__enter__()
+            cls._restart_celery_worker()
 
         if os.path.isdir(SCR_DIR):
             rmtree(SCR_DIR)
@@ -264,6 +279,9 @@ def setUp(self):
         self.current_order_id = None
         self.order_ids_before_launch = None
 
+        if not IS_CLOUD:
+            self.__class__._restart_celery_worker()
+
         self.user = User.objects.create_user(
             email=self.email,
             password=self.password,
@@ -1530,10 +1548,19 @@ def click_latest_calc(self):
         assert self.is_on_page_calculations()
         assert self.get_number_calc_orders() > 0
 
-        calculations = self.get_calc_orders()
-        self.current_order_id = calculations[0].get_attribute("id")
+        target = self._get_target_calc_order()
+        if target is None:
+            target_order_id = self._get_target_order_id()
+            if target_order_id is not None:
+                self.current_order_id = f"order_{target_order_id}"
+                self.order_ids_before_launch = None
+                self.driver.get(f"{self.live_server_url}/link_order/{target_order_id}")
+                self.wait_for_ajax()
+                return
+            target = self.get_calc_orders()[0]
+        self.current_order_id = target.get_attribute("id")
         self.order_ids_before_launch = None
-        link = calculations[0].find_element(By.CSS_SELECTOR, "a[href^='/link_order/']")
+        link = target.find_element(By.CSS_SELECTOR, "a[href^='/link_order/']")
         self.driver.get(link.get_attribute("href"))
         self.wait_for_ajax()
 
@@ -1551,32 +1578,58 @@ def details_latest_order(self):
         assert self.is_on_page_calculations()
         assert self.get_number_calc_orders() > 0
 
-        calculations = self.get_calc_orders()
-        self.current_order_id = calculations[0].get_attribute("id")
+        target = self._get_target_calc_order()
+        if target is None:
+            target_order_id = self._get_target_order_id()
+            if target_order_id is not None:
+                self.current_order_id = f"order_{target_order_id}"
+                self.order_ids_before_launch = None
+                self.driver.get(
+                    f"{self.live_server_url}/calculationorder/{target_order_id}"
+                )
+                self.wait_for_ajax()
+                return
+            target = self.get_calc_orders()[0]
+        self.current_order_id = target.get_attribute("id")
         self.order_ids_before_launch = None
 
-        icon = calculations[0].find_element(By.CLASS_NAME, "fa-list")
+        icon = target.find_element(By.CLASS_NAME, "fa-list")
         link = icon.find_element(By.XPATH, "./ancestor::a[1]")
         self.driver.get(link.get_attribute("href"))
         self.wait_for_ajax()
 
+    def _get_target_order_id(self):
+        if self.current_order_id is not None:
+            return self.current_order_id.removeprefix("order_")
+
+        if self.order_ids_before_launch is None:
+            return None
+
+        close_old_connections()
+        current_ids = {
+            f"order_{order_id}"
+            for order_id in CalculationOrder.objects.values_list("id", flat=True)
+        }
+        new_ids = current_ids - self.order_ids_before_launch
+        if not new_ids:
+            return None
+
+        target_id = sorted(new_ids)[-1]
+        self.current_order_id = target_id
+        self.order_ids_before_launch = None
+        return target_id.removeprefix("order_")
+
     def _get_target_calc_order(self):
         calculations = self.get_calc_orders()
-        if self.current_order_id is None and self.order_ids_before_launch is not None:
-            for calc in calculations:
-                calc_id = calc.get_attribute("id")
-                if calc_id not in self.order_ids_before_launch:
-                    self.current_order_id = calc_id
-                    self.order_ids_before_launch = None
-                    return calc
-            return None
-        if self.current_order_id is None:
+        target_order_id = self._get_target_order_id()
+        if target_order_id is None:
             return calculations[0]
 
+        target_dom_id = f"order_{target_order_id}"
         for calc in calculations:
-            if calc.get_attribute("id") == self.current_order_id:
+            if calc.get_attribute("id") == target_dom_id:
                 return calc
-        return calculations[0]
+        return None
 
     def _refresh_calculations_page(self):
         self.driver.refresh()
@@ -1713,6 +1766,18 @@ def wait_latest_calc_done(self, timeout):
                 self._refresh_calculations_page()
                 return
 
+            target_order_id = self._get_target_order_id()
+            if target_order_id is not None:
+                close_old_connections()
+                backend_status = (
+                    CalculationOrder.objects.filter(id=target_order_id)
+                    .values_list("cached_status", flat=True)
+                    .first()
+                )
+                if backend_status in [2, 3]:
+                    self._refresh_calculations_page()
+                    return
+
             time.sleep(0.5)
             self._refresh_calculations_page()
         raise Exception("Calculation did not finish")

frontend/models.py

@@ -238,12 +238,25 @@ def bill_time(self, time):
 
         return user
 
+    @classmethod
+    def from_db(cls, db, field_names, values):
+        instance = super().from_db(db, field_names, values)
+        if "unseen_calculations" in field_names:
+            instance._loaded_unseen_calculations = instance.unseen_calculations
+        else:
+            instance._loaded_unseen_calculations = None
+        return instance
+
     def save(self, *args, **kwargs):
         update_fields = kwargs.get("update_fields")
-        should_preserve_unseen = (
-            self.pk is not None
-            and update_fields is not None
-            and "unseen_calculations" not in update_fields
+        should_preserve_unseen = self.pk is not None and (
+            (update_fields is not None and "unseen_calculations" not in update_fields)
+            or (
+                update_fields is None
+                and hasattr(self, "_loaded_unseen_calculations")
+                and self._loaded_unseen_calculations is not None
+                and self.unseen_calculations == self._loaded_unseen_calculations
+            )
         )
         if should_preserve_unseen:
             current_unseen = (
@@ -255,6 +268,7 @@ def save(self, *args, **kwargs):
                 self.unseen_calculations = current_unseen
 
         super().save(*args, **kwargs)
+        self._loaded_unseen_calculations = self.unseen_calculations
 
     @property
     def is_PI(self):
@@ -1401,7 +1415,7 @@ def _get_label(self):
     def molecule_name(self):
         if self._molecule_name == "":
             self._molecule_name = self._get_molecule_name()
-            self.save()
+            self.save(update_fields=["_molecule_name"])
         return self._molecule_name
 
     @property
@@ -1411,7 +1425,7 @@ def step_name(self):
                 self._step_name = "Unknown"
             else:
                 self._step_name = self.step.name
-            self.save()
+            self.save(update_fields=["_step_name"])
         return self._step_name
 
     @property
@@ -1421,7 +1435,7 @@ def project_name(self):
                 self._project_name = "Unknown"
             else:
                 self._project_name = self.project.name
-            self.save()
+            self.save(update_fields=["_project_name"])
         return self._project_name
 
     def _get_molecule_name(self):