Add more suggest next sub ses tests (#597)

* Fix get next sub / ses on TUI and extend AWS, SSH, GDrive tests.

* Also remove aws folder.

* Apply PR feedback and refactor tests.

* Move check_next_sub_ses_in_tui to tui base.

* Update tests/tests_transfers/ssh/test_ssh_suggest_next.py

Co-authored-by: Shrey Singh <96627769+cs7-shrey@users.noreply.github.com>

* Update tests/tests_transfers/aws/test_aws_suggest_next.py

Co-authored-by: Shrey Singh <96627769+cs7-shrey@users.noreply.github.com>

---------

Co-authored-by: Shrey Singh <96627769+cs7-shrey@users.noreply.github.com>

Author: JoeZiminski

datashuttle/tui/tabs/create_folders.py

@@ -338,10 +338,9 @@ def suggest_next_sub_ses(
             in canonical_configs.get_connection_methods_list()
         )
 
-        if (
-            include_central
-            and self.interface.project.cfg["connection_method"] == "ssh"
-        ):
+        if include_central and self.interface.project.cfg[
+            "connection_method"
+        ] in ["aws", "gdrive", "ssh"]:
             self.searching_central_popup_widget = (
                 SearchingCentralForNextSubSesPopup(prefix)
             )

tests/tests_transfers/aws/test_aws_suggest_next.py

@@ -0,0 +1,41 @@
+import pytest
+
+from datashuttle.utils import rclone
+
+from ... import test_utils
+from ...tests_tui.tui_base import TuiBase
+from ..base_transfer import BaseTransfer
+from . import aws_test_utils
+
+
+@pytest.mark.skipif(
+    not aws_test_utils.has_aws_environment_variables(),
+    reason="AWS set up environment variables must be set.",
+)
+class TestAWSSuggestNext(BaseTransfer, TuiBase):
+    @pytest.fixture(
+        scope="function",
+    )
+    def aws_setup(self, setup_project_paths):
+        """
+        Setup pathtable and project for AWS transfer tests.
+        """
+        project = test_utils.make_project(setup_project_paths["project_name"])
+        aws_test_utils.setup_project_for_aws(project)
+        aws_test_utils.setup_aws_connection(project)
+
+        yield project
+
+        rclone.call_rclone(
+            f"purge central_{project.project_name}_gdrive:{project.cfg['central_path'].parent}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_aws_suggest_next_sub_ses(
+        self,
+        aws_setup,
+    ):
+        """ """
+        project = aws_setup
+
+        await self.check_next_sub_ses_in_tui(project)

tests/tests_transfers/gdrive/test_gdrive_suggest_next.py

@@ -0,0 +1,43 @@
+import pytest
+
+from datashuttle.utils import rclone
+
+from ... import test_utils
+from ...tests_tui.tui_base import TuiBase
+from ..base_transfer import BaseTransfer
+from . import gdrive_test_utils
+
+
+@pytest.mark.skipif(
+    not gdrive_test_utils.has_gdrive_environment_variables(),
+    reason="Google Drive set up environment variables must be set.",
+)
+class TestGDriveSuggestNext(BaseTransfer, TuiBase):
+    @pytest.fixture(
+        scope="function",
+    )
+    def gdrive_setup(self, setup_project_paths):
+        """
+        Setup pathtable and project for GDrive transfer tests.
+        """
+        project = test_utils.make_project(setup_project_paths["project_name"])
+        gdrive_test_utils.setup_project_for_gdrive(
+            project,
+        )
+        gdrive_test_utils.setup_gdrive_connection(project)
+
+        yield project
+
+        rclone.call_rclone(
+            f"purge central_{project.project_name}_gdrive:{project.cfg['central_path'].parent}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_gdrive_suggest_next_sub_ses(
+        self,
+        gdrive_setup,
+    ):
+        """ """
+        project = gdrive_setup
+
+        await self.check_next_sub_ses_in_tui(project)

tests/tests_transfers/ssh/base_ssh.py

@@ -25,7 +25,7 @@ class BaseSSHTransfer(BaseTransfer):
     @pytest.fixture(
         scope="class",
     )
-    def setup_ssh_container(self):
+    def setup_ssh_container_fixture(self):
         """
         Set up the Dockerfile container for SSH tests and
         delete it on teardown.

tests/tests_transfers/ssh/test_ssh_setup.py

@@ -23,7 +23,7 @@
 )
 class TestSSH(BaseSSHTransfer):
     @pytest.fixture(scope="function")
-    def project(test, tmp_path, setup_ssh_container):
+    def project(test, tmp_path, setup_ssh_container_fixture):
         """Set up a project with configs for SSH into
         the test Dockerfile image.
         """

tests/tests_transfers/ssh/test_ssh_suggest_next.py

@@ -0,0 +1,45 @@
+import platform
+
+import pytest
+
+from ... import test_utils
+from ...tests_tui.tui_base import TuiBase
+from . import ssh_test_utils
+from .base_ssh import BaseSSHTransfer
+
+TEST_SSH = ssh_test_utils.docker_is_running()
+
+
+@pytest.mark.skipif(
+    platform.system == "Darwin", reason="Docker set up is not robust on macOS."
+)
+@pytest.mark.skipif(
+    not TEST_SSH,
+    reason="SSH tests are not run as docker is either not installed, "
+    "running or current user is not in the docker group.",
+)
+class TestSSHDriveSuggestNext(BaseSSHTransfer, TuiBase):
+    @pytest.fixture(
+        scope="function",
+    )
+    def ssh_setup(self, setup_project_paths, setup_ssh_container_fixture):
+        """
+        Setup pathtable and project for SSH transfer tests.
+        """
+        project = test_utils.make_project(setup_project_paths["project_name"])
+        ssh_test_utils.setup_project_for_ssh(
+            project,
+        )
+        ssh_test_utils.setup_ssh_connection(project)
+
+        yield project
+
+    @pytest.mark.asyncio
+    async def test_ssh_suggest_next_sub_ses(
+        self,
+        ssh_setup,
+    ):
+        """ """
+        project = ssh_setup
+
+        await self.check_next_sub_ses_in_tui(project)

tests/tests_transfers/ssh/test_ssh_transfer.py

@@ -21,7 +21,7 @@ class TestSSHTransfer(BaseSSHTransfer):
     @pytest.fixture(
         scope="class",
     )
-    def ssh_setup(self, pathtable_and_project, setup_ssh_container):
+    def ssh_setup(self, pathtable_and_project, setup_ssh_container_fixture):
         """
         After initial project setup (in `pathtable_and_project`)
         setup a container and the project's SSH connection to the container.

tests/tests_tui/aws/__init__.py

@@ -762,20 +762,3 @@ async def create_folders_and_check_output(
             sessions=sessions,
             folder_used=folder_used,
         )
-
-    async def double_click_input(self, pilot, sub_or_ses, control=False):
-        """Helper function to double click input to suggest next sub or ses.
-
-        Because this function is performed in separate asyncio task, this was a little
-        brittle in the CI tests leading to random errors. The below
-        combination of awaiting the test, then pausing, stopped the errors.
-        """
-        expand_name = "session" if sub_or_ses == "ses" else "subject"
-
-        await self.double_click(
-            pilot, f"#create_folders_{expand_name}_input", control=control
-        )
-        await test_utils.await_task_by_name_if_present(
-            f"suggest_next_{sub_or_ses}_async_task"
-        )
-        await pilot.pause(0.5)

tests/tests_tui/test_tui_create_folders.py

@@ -1,7 +1,10 @@
+import shutil
+
 import pytest_asyncio
 from textual.widgets._tabbed_content import ContentTab
 
 from datashuttle.configs import canonical_configs
+from datashuttle.tui.app import TuiApp
 from datashuttle.tui.screens.project_manager import ProjectManagerScreen
 from datashuttle.tui.screens.project_selector import ProjectSelectorScreen
 
@@ -23,6 +26,9 @@ def tui_size(self):
         """
         return (500, 500)
 
+    # Fixtures
+    # ---------------------------------------------------------------------------------
+
     @pytest_asyncio.fixture(scope="function")
     async def empty_project_paths(self, tmp_path_factory, monkeypatch):
         """Get the paths and project name for a non-existent (i.e. not
@@ -64,6 +70,9 @@ async def setup_project_paths(self, empty_project_paths):
 
         return empty_project_paths
 
+    # Helper Functions
+    # ---------------------------------------------------------------------------------
+
     def monkeypatch_print(self, _monkeypatch):
         """Calls to `print` in datashuttle crash the TUI in the
         test environment. I am not sure why. Get around this
@@ -240,3 +249,84 @@ async def click_and_await_transfer(self, pilot):
             await transfer_task
 
         await self.close_messagebox(pilot)
+
+    async def double_click_input(self, pilot, sub_or_ses, control=False):
+        """Helper function to double click input to suggest next sub or ses.
+
+        Because this function is performed in separate asyncio task, this was a little
+        brittle in the CI tests leading to random errors. The below
+        combination of awaiting the test, then pausing, stopped the errors.
+        """
+        expand_name = "session" if sub_or_ses == "ses" else "subject"
+
+        await self.double_click(
+            pilot, f"#create_folders_{expand_name}_input", control=control
+        )
+        await test_utils.await_task_by_name_if_present(
+            f"suggest_next_{sub_or_ses}_async_task"
+        )
+        await pilot.pause(0.5)
+
+    # Shared checks
+    # ---------------------------------------------------------------------------------
+
+    async def check_next_sub_ses_in_tui(self, project):
+        """A central function for testing next sub / ses in the TUI.
+
+        This test is shared between ssh, aws and gdrive tests that
+        use the same logic to test the get next sub / ses functionality.
+
+        First, sub / ses folders are created in the project and uploaded
+        centrally. Then, the folders are removed  from the local path.
+        In this way, we can be use that `include_central` (which searches
+        the remote for the next sub / ses) is behaving correctly and not just
+        reading the local path.
+
+        Because sub-001 is created, the suggested sub we expect is sub-002.
+        Because ses-002 is created in sub-001, the suggested ses we expect is ses-003.
+        """
+        test_utils.make_local_folders_with_files_in(
+            project, "rawdata", "sub-001", ["ses-001", "ses-002"]
+        )
+        project.upload_entire_project()
+
+        shutil.rmtree(project.get_local_path())
+
+        app = TuiApp()
+        async with app.run_test(size=self.tui_size()) as pilot:
+            await self.check_and_click_onto_existing_project(
+                pilot, project.project_name
+            )
+
+            # Turn on the central checkbox
+            await self.scroll_to_click_pause(
+                pilot, "#create_folders_settings_button"
+            )
+            await self.scroll_to_click_pause(
+                pilot, "#suggest_next_sub_ses_central_checkbox"
+            )
+            await self.scroll_to_click_pause(
+                pilot, "#create_folders_settings_close_button"
+            )
+
+            await self.fill_input(
+                pilot, "#create_folders_subject_input", "sub-001"
+            )
+
+            await self.double_click_input(pilot, "ses")
+
+            assert (
+                pilot.app.screen.query_one(
+                    "#create_folders_session_input"
+                ).value
+                == "ses-003"
+            )
+
+            await self.double_click_input(pilot, "sub")
+
+            assert (
+                pilot.app.screen.query_one(
+                    "#create_folders_subject_input"
+                ).value
+                == "sub-002"
+            )