Merge pull request #567 from cs7-shrey/gdrive_aws_tests

ADD: Tests for Google Drive and AWS Connection Methods

Author: cs7-shrey

.github/workflows/code_test_and_deploy.yml

@@ -60,6 +60,16 @@ jobs:
         # run SSH tests only on Linux because Windows and macOS
         # are already run within a virtual container and so cannot
         # run Linux containers because nested containerisation is disabled.
+      - name: Setup Google Drive and AWS Secrets
+        run: |
+          echo "GDRIVE_CLIENT_ID=${{secrets.GDRIVE_CLIENT_ID}}" >> $GITHUB_ENV
+          echo "GDRIVE_CLIENT_SECRET=${{secrets.GDRIVE_CLIENT_SECRET}}" >> $GITHUB_ENV
+          echo "GDRIVE_ROOT_FOLDER_ID=${{secrets.GDRIVE_ROOT_FOLDER_ID}}" >> $GITHUB_ENV
+          echo "GDRIVE_CONFIG_TOKEN=${{secrets.GDRIVE_CONFIG_TOKEN}}" >> $GITHUB_ENV
+          echo "AWS_ACCESS_KEY_ID=${{secrets.AWS_ACCESS_KEY_ID}}" >> $GITHUB_ENV
+          echo "AWS_REGION=${{secrets.AWS_REGION}}" >> $GITHUB_ENV
+          echo "AWS_BUCKET_NAME=${{secrets.AWS_BUCKET_NAME}}" >> $GITHUB_ENV
+          echo "AWS_SECRET_ACCESS_KEY=${{secrets.AWS_SECRET_ACCESS_KEY}}" >> $GITHUB_ENV
       - name: Test SSH (Linux only)
         if: runner.os == 'Linux'
         run: |

datashuttle/tui/screens/setup_gdrive.py

@@ -191,7 +191,8 @@ def open_browser_and_setup_gdrive_connection(
         asyncio.create_task(
             self.setup_gdrive_connection_and_update_ui(
                 gdrive_client_secret=gdrive_client_secret
-            )
+            ),
+            name="setup_gdrive_connection_with_browser_task",
         )
 
     def prompt_user_for_config_token(self) -> None:
@@ -227,7 +228,8 @@ def setup_gdrive_connection_using_config_token(
             self.setup_gdrive_connection_and_update_ui(
                 gdrive_client_secret=gdrive_client_secret,
                 config_token=config_token,
-            )
+            ),
+            name="setup_gdrive_connection_without_browser_task",
         )
 
     async def setup_gdrive_connection_and_update_ui(

datashuttle/tui/shared/configs_content.py

@@ -129,7 +129,7 @@ def compose(self) -> ComposeResult:
             ClickableInput(
                 self.parent_class.mainwindow,
                 placeholder="Google Drive Root Folder ID",
-                id="configs_gdrive_root_folder_id",
+                id="configs_gdrive_root_folder_id_input",
             ),
             Label("Client ID (Optional)", id="configs_gdrive_client_id_label"),
             ClickableInput(
@@ -301,7 +301,7 @@ def on_mount(self) -> None:
             "#configs_central_host_username_input",
             "#configs_central_host_id_input",
             "#configs_gdrive_client_id_input",
-            "#configs_gdrive_root_folder_id",
+            "#configs_gdrive_root_folder_id_input",
         ]:
             self.query_one(id).tooltip = get_tooltip(id)
 
@@ -721,7 +721,7 @@ def fill_inputs_with_project_configs(self) -> None:
         input.value = value
 
         # Google Drive Root Folder ID
-        input = self.query_one("#configs_gdrive_root_folder_id")
+        input = self.query_one("#configs_gdrive_root_folder_id_input")
         value = (
             ""
             if cfg_to_load["gdrive_root_folder_id"] is None
@@ -897,7 +897,7 @@ def get_datashuttle_inputs_from_widgets(self) -> Dict:
 
             cfg_kwargs["gdrive_root_folder_id"] = (
                 self.get_config_value_from_input_value(
-                    "#configs_gdrive_root_folder_id"
+                    "#configs_gdrive_root_folder_id_input"
                 )
             )
 

datashuttle/tui/tooltips.py

@@ -84,7 +84,7 @@ def get_tooltip(id: str) -> str:
             "Can be left empty to use rclone's default client (slower)"
         )
 
-    elif id == "#configs_gdrive_root_folder_id":
+    elif id == "#configs_gdrive_root_folder_id_input":
         tooltip = (
             "The Google Drive root folder ID to use for transfer.\n\n"
             "It can be obtained by navigating to the folder in Google Drive and copying the ID from the URL.\n\n"

tests/test_utils.py

@@ -1,6 +1,7 @@
 import asyncio
 import copy
 import glob
+import json
 import logging
 import os
 import pathlib
@@ -414,6 +415,49 @@ def check_config_file(config_path, *kwargs):
             assert value == config_yaml[name], f"{name}"
 
 
+# -----------------------------------------------------------------------------
+# Search
+# -----------------------------------------------------------------------------
+
+
+def recursive_search_central(project: DataShuttle):
+    """
+    A convenience function to recursively search a
+    project for files on remote folders using `rclone
+    lsjson`
+    """
+    all_filenames: list[str] = []
+    rclone_recursive_search(
+        project,
+        (project.cfg["central_path"] / "rawdata").as_posix(),
+        all_filenames,
+    )
+
+    return all_filenames
+
+
+def rclone_recursive_search(
+    project: DataShuttle, path_: str, all_filenames: list[str]
+):
+    output = rclone.call_rclone(
+        f"lsjson {project.cfg.get_rclone_config_name()}:{path_}", pipe_std=True
+    )
+
+    files_or_folders = json.loads(output.stdout)
+
+    for file_or_folder in files_or_folders:
+        name = file_or_folder["Name"]
+        is_dir = file_or_folder.get("IsDir", False)
+        if is_dir:
+            rclone_recursive_search(
+                project,
+                path_ + "/" + name,
+                all_filenames,
+            )
+        else:
+            all_filenames.append(path_ + "/" + name)
+
+
 # -----------------------------------------------------------------------------
 # Test Helpers
 # -----------------------------------------------------------------------------

tests/tests_transfers/aws/__init__.py

@@ -0,0 +1,30 @@
+import copy
+import os
+
+from datashuttle import DataShuttle
+from datashuttle.utils import aws
+
+
+def setup_project_for_aws(project: DataShuttle):
+    aws_bucket_name = os.environ["AWS_BUCKET_NAME"]
+    project.update_config_file(
+        connection_method="aws",
+        central_path=f"{aws_bucket_name}/main/{project.project_name}",
+        aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
+        aws_region=os.environ["AWS_REGION"],
+    )
+
+
+def setup_aws_connection(project: DataShuttle):
+    """
+    Convenience function to set up the AWS connection by
+    mocking the `aws.get_aws_secret_access_key` function.
+    """
+    original_get_secret = copy.deepcopy(aws.get_aws_secret_access_key)
+    aws.get_aws_secret_access_key = lambda *args, **kwargs: os.environ[
+        "AWS_SECRET_ACCESS_KEY"
+    ]
+
+    project.setup_aws_connection()
+
+    aws.get_aws_secret_access_key = original_get_secret

tests/tests_transfers/aws/aws_test_utils.py

@@ -0,0 +1,50 @@
+import pytest
+
+from datashuttle.utils import rclone
+
+from ..base_transfer import BaseTransfer
+from . import aws_test_utils
+
+
+class TestAwsTransfer(BaseTransfer):
+    @pytest.fixture(
+        scope="class",
+    )
+    def aws_setup(self, pathtable_and_project):
+        """
+        Setup pathtable and project for AWS transfer tests.
+        """
+        pathtable, project = pathtable_and_project
+
+        aws_test_utils.setup_project_for_aws(project)
+        aws_test_utils.setup_aws_connection(project)
+
+        project.upload_rawdata()
+
+        yield [pathtable, project]
+
+        rclone.call_rclone(
+            f"purge central_{project.project_name}_aws:{project.cfg['central_path']}"
+        )
+
+    @pytest.mark.parametrize(
+        "sub_names", [["all"], ["all_non_sub", "sub-002"]]
+    )
+    @pytest.mark.parametrize(
+        "ses_names", [["all"], ["ses-002_random-key"], ["all_non_ses"]]
+    )
+    @pytest.mark.parametrize(
+        "datatype", [["all"], ["anat", "all_non_datatype"]]
+    )
+    def test_combinations_aws_transfer(
+        self,
+        aws_setup,
+        sub_names,
+        ses_names,
+        datatype,
+    ):
+        pathtable, project = aws_setup
+
+        self.run_and_check_transfers(
+            project, pathtable, sub_names, ses_names, datatype
+        )

tests/tests_transfers/aws/test_aws_transfer.py

@@ -1,11 +1,14 @@
 """ """
 
 import copy
+import shutil
 from pathlib import Path
 
 import pandas as pd
 import pytest
 
+from datashuttle.utils import rclone
+
 from .. import test_utils
 from ..base import BaseTest
 from .file_conflicts_pathtable import get_pathtable
@@ -159,3 +162,81 @@ def create_all_pathtable_files(self, pathtable):
 
     def central_from_local(self, path_):
         return Path(str(copy.copy(path_)).replace("local", "central"))
+
+    # -------------------------------------------------------------------------
+    # Transfer
+    # -------------------------------------------------------------------------
+
+    def run_and_check_transfers(
+        self, project, pathtable, sub_names, ses_names, datatype
+    ):
+        true_central_path = project.cfg["central_path"]
+        tmp_central_path = (
+            project.cfg["central_path"] / "tmp" / project.project_name
+        )
+        self.remake_logging_path(project)
+
+        project.update_config_file(central_path=tmp_central_path)
+
+        project.upload_custom(
+            "rawdata", sub_names, ses_names, datatype, init_log=False
+        )
+
+        expected_transferred_paths = self.get_expected_transferred_paths(
+            pathtable, sub_names, ses_names, datatype
+        )
+
+        transferred_files = test_utils.recursive_search_central(project)
+        paths_to_transferred_files = self.remove_path_before_rawdata(
+            transferred_files
+        )
+
+        assert sorted(paths_to_transferred_files) == sorted(
+            expected_transferred_paths
+        )
+
+        # Now, move data from the central path where the project is
+        # setup, to a temp local folder to test download.
+        true_local_path = project.cfg["local_path"]
+        tmp_local_path = (
+            project.cfg["local_path"] / "tmp" / project.project_name
+        )
+        tmp_local_path.mkdir(exist_ok=True, parents=True)
+
+        project.update_config_file(local_path=tmp_local_path)
+        project.update_config_file(central_path=true_central_path)
+
+        project.download_custom(
+            "rawdata", sub_names, ses_names, datatype, init_log=False
+        )
+
+        # Find the transferred paths, tidy them up
+        # and check expected paths were transferred.
+        all_transferred = list((tmp_local_path / "rawdata").glob("**/*"))
+        all_transferred = [
+            path_ for path_ in all_transferred if path_.is_file()
+        ]
+
+        paths_to_transferred_files = self.remove_path_before_rawdata(
+            all_transferred
+        )
+
+        assert sorted(paths_to_transferred_files) == sorted(
+            expected_transferred_paths
+        )
+
+        rclone.call_rclone(
+            f"purge {project.cfg.get_rclone_config_name()}:{tmp_central_path.as_posix()}"
+        )
+
+        shutil.rmtree(tmp_local_path)
+
+        self.remake_logging_path(project)
+        project.update_config_file(local_path=true_local_path)
+
+    def remake_logging_path(self, project):
+        """
+        Need to do this to compensate for switching
+        local_path location in the test environment.
+        """
+        project.get_logging_path().mkdir(parents=True, exist_ok=True)