refactor: apply code review changes

Author: shrey

datashuttle/tui/screens/modal_dialogs.py

@@ -138,6 +138,13 @@ async def handle_transfer_and_update_ui_when_complete(self) -> None:
 
 
 class SearchingCentralForNextSubSesPopup(ModalScreen):
+    """
+    A popup to show message and a loading indicator when awaiting search next sub/ses across
+    the folders present in both local and central machines. This search happens in a separate
+    thread so as to allow TUI to display the loading indicate without freezing.
+
+    Only displayed when the `include_central` flag is checked and the connection method is "ssh".
+    """
 
     def __init__(self, sub_or_ses: Prefix) -> None:
         super().__init__()

datashuttle/tui/tabs/create_folders.py

@@ -305,11 +305,10 @@ def suggest_next_sub_ses(
             include_central
             and self.interface.project.cfg["connection_method"] == "ssh"
         ):
-            searching_central_popup = SearchingCentralForNextSubSesPopup(
-                prefix
+            self.searching_central_popup_widget = (
+                SearchingCentralForNextSubSesPopup(prefix)
             )
-            self.searching_central_popup_widget = searching_central_popup
-            self.mainwindow.push_screen(searching_central_popup)
+            self.mainwindow.push_screen(self.searching_central_popup_widget)
 
         asyncio.create_task(
             self.fill_suggestion_and_dismiss_popup(

tests/test_utils.py

@@ -695,3 +695,8 @@ def get_task_by_name(name):
         None,
     )
     return target_task
+
+
+async def await_task_by_name_if_present(name: str) -> None:
+    if task := get_task_by_name(name):
+        await task

tests/tests_tui/test_tui_create_folders.py

@@ -445,10 +445,9 @@ async def test_name_template_next_sub_or_ses_and_validation(
             await self.double_click(
                 pilot, "#create_folders_subject_input", control=False
             )
-            if suggest_sub_task := test_utils.get_task_by_name(
+            await test_utils.await_task_by_name_if_present(
                 "suggest_next_sub_async_task"
-            ):
-                await suggest_sub_task
+            )
             assert (
                 pilot.app.screen.query_one(
                     "#create_folders_subject_input"
@@ -462,10 +461,9 @@ async def test_name_template_next_sub_or_ses_and_validation(
             await self.double_click(
                 pilot, "#create_folders_session_input", control=False
             )
-            if suggest_ses_task := test_utils.get_task_by_name(
+            await test_utils.await_task_by_name_if_present(
                 "suggest_next_ses_async_task"
-            ):
-                await suggest_ses_task
+            )
             assert (
                 pilot.app.screen.query_one(
                     "#create_folders_session_input"
@@ -512,10 +510,9 @@ async def test_get_next_sub_and_ses_no_template(self, setup_project_paths):
 
             # Double click without CTRL modifier key.
             await self.double_click(pilot, "#create_folders_subject_input")
-            if suggest_sub_task := test_utils.get_task_by_name(
+            await test_utils.await_task_by_name_if_present(
                 "suggest_next_sub_async_task"
-            ):
-                await suggest_sub_task
+            )
             assert (
                 pilot.app.screen.query_one(
                     "#create_folders_subject_input"
@@ -527,10 +524,9 @@ async def test_get_next_sub_and_ses_no_template(self, setup_project_paths):
                 pilot, "#create_folders_subject_input", "sub-001"
             )
             await self.double_click(pilot, "#create_folders_session_input")
-            if suggest_ses_task := test_utils.get_task_by_name(
+            await test_utils.await_task_by_name_if_present(
                 "suggest_next_ses_async_task"
-            ):
-                await suggest_ses_task
+            )
             assert (
                 pilot.app.screen.query_one(
                     "#create_folders_session_input"
@@ -544,6 +540,11 @@ async def test_get_next_sub_and_ses_no_template(self, setup_project_paths):
     async def test_get_next_sub_and_ses_central_no_template(
         self, setup_project_paths, mocker
     ):
+        """
+        Test getting the next subject / session with the include_central option. Check the
+        checkbox widget that turns the setting on. Trigger a get next subject / session and mock
+        the underlying datashuttle function to ensure include_central is properly called.
+        """
         tmp_config_path, tmp_path, project_name = setup_project_paths.values()
 
         app = TuiApp()
@@ -571,13 +572,11 @@ async def test_get_next_sub_and_ses_central_no_template(
                 pilot.app.screen.interface.project, "get_next_ses"
             )
 
-            #  # Check subject suggestion called mocked function correctly
+            # Check subject suggestion called mocked function correctly
             await self.double_click(pilot, "#create_folders_subject_input")
-
-            if suggest_sub_task := test_utils.get_task_by_name(
+            await test_utils.await_task_by_name_if_present(
                 "suggest_next_sub_async_task"
-            ):
-                await suggest_sub_task
+            )
 
             spy_get_next_sub.assert_called_with(
                 "rawdata", return_with_prefix=True, include_central=True
@@ -589,10 +588,9 @@ async def test_get_next_sub_and_ses_central_no_template(
             )
             await self.double_click(pilot, "#create_folders_session_input")
 
-            if suggest_ses_task := test_utils.get_task_by_name(
+            await test_utils.await_task_by_name_if_present(
                 "suggest_next_ses_async_task"
-            ):
-                await suggest_ses_task
+            )
 
             spy_get_next_ses.assert_called_with(
                 "rawdata",

tests/tests_tui/test_tui_widgets_and_defaults.py

@@ -931,6 +931,11 @@ async def check_top_folder_select(
     async def test_search_central_for_suggestion_settings(
         self, setup_project_paths
     ):
+        """
+        Check the settings for the checkbox that selects include_central when
+        getting the next subject or session in the 'Create' tab and ensure that
+        the underlying settings are changed.
+        """
         tmp_config_path, tmp_path, project_name = setup_project_paths.values()
 
         app = TuiApp()