refactor: build the transport at init and reach Google Drive only through its components

packages/syft-enclave/src/syft_enclaves/client.py

@@ -15,6 +15,7 @@
 from syft_rds import SyftRDSClient, SyftRDSClientConfig
 from syft_rds.waiting import JobCheck, JobStatusArg, wait_for_job
 
+from syft.sync.transport.protocols import Transport
 from syft.sync.encryption_policy import EncryptionPolicy
 from syft.sync.peers.peer import Peer
 from syft.sync.peers.peer_list import PeerList
@@ -45,7 +46,6 @@
     create_clients,
     create_configs,
     setup_callbacks,
-    setup_connections,
     wire_peers,
     write_versions,
 )
@@ -783,15 +783,18 @@ def from_config(
         cls,
         config: SyftRDSClientConfig,
         data_owners: list[str] | None = None,
+        transport: Transport | None = None,
     ) -> "SyftEnclaveClient":
         """Build a SyftEnclaveClient from an RDS config with a wrapped job_client.
 
         Encryption rides along on ``config.sync`` via
         ``peer_manager_config.encryption_policy``; ``SyftboxManager.from_config``
         (invoked inside ``SyftRDSClient.from_config``) loads/generates this
         datasite's keys from its per-datasite key file, so no extra step is needed.
+        ``transport`` replaces the one ``config.sync`` describes; tests pass
+        one over a mock Drive.
         """
-        rds = SyftRDSClient.from_config(config)
+        rds = SyftRDSClient.from_config(config, transport)
         rds.job_client = EnclaveJobClient(rds.job_client)
 
         sync_engine = rds.sync_engine
@@ -885,12 +888,10 @@ def quad_with_mock_drive_service_connection(
         enclave, do1, do2, ds = clients
         enclave.data_owners = [do1.email, do2.email]
         # The setup helpers manipulate the generic sync engine directly
-        # (connections, file-watcher callbacks, version files, encryption,
-        # peering). Each client's sync engine lives at ``rds.sync_engine``; the
+        # (file-watcher callbacks, version files, encryption, peering). Each client's sync engine lives at ``rds.sync_engine``; the
         # RDS-owned job/dataset managers react via the peer-lifecycle callbacks
         # registered on that same engine during SyftRDSClient.__init__.
         managers = tuple(c._rds.sync_engine for c in clients)
-        setup_connections(managers)
         setup_callbacks(managers)
         write_versions(managers)
         wire_peers(managers)

packages/syft-enclave/src/syft_enclaves/utils.py

@@ -4,6 +4,8 @@
     MockDriveService,
 )
 from syft.sync.connections.drive.gdrive_transport import GDriveConnection
+from syft.sync.syftbox_manager import SyftboxManager
+from syft.sync.transport.protocols import Transport
 
 
 def create_configs(
@@ -49,18 +51,23 @@ def create_configs(
 
 
 def create_clients(configs) -> tuple:
+    """One client per config, all on one mock Drive."""
     from syft_enclaves.client import SyftEnclaveClient
 
-    return tuple(SyftEnclaveClient.from_config(c) for c in configs)
-
-
-def setup_connections(managers: tuple):
     backing_store = MockDriveBackingStore()
+    return tuple(
+        SyftEnclaveClient.from_config(
+            c, transport=_mock_transport(backing_store, c.sync.email)
+        )
+        for c in configs
+    )
 
-    for manager in managers:
-        service = MockDriveService(backing_store, manager.email)
-        connection = GDriveConnection.from_service(manager.email, service)
-        manager._add_connection(connection)
+
+def _mock_transport(backing_store: MockDriveBackingStore, email: str) -> Transport:
+    service = MockDriveService(backing_store, email)
+    return SyftboxManager._drive_transport(
+        GDriveConnection.from_service(email, service)
+    )
 
 
 def setup_callbacks(managers: tuple):

packages/syft-rds/src/syft_rds/client.py

@@ -13,6 +13,7 @@
 from pydantic import BaseModel, ConfigDict
 from syft_migration import ProtocolSchema
 
+from syft.sync.transport.protocols import Transport
 from syft.sync.syftbox_manager import SyftboxManager
 from syft.sync.peers.peer import Peer
 from syft.sync.utils.waiting import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIMEOUT
@@ -148,8 +149,14 @@ def model_post_init(self, __context: Any) -> None:
             )
 
     @classmethod
-    def from_config(cls, config: "SyftRDSClientConfig") -> "SyftRDSClient":
-        sync_engine = SyftboxManager.from_config(config.sync)
+    def from_config(
+        cls,
+        config: "SyftRDSClientConfig",
+        transport: Transport | None = None,
+    ) -> "SyftRDSClient":
+        """``transport`` replaces the one ``config.sync`` describes; tests pass
+        one over a mock Drive (see ``SyftboxManager.from_config``)."""
+        sync_engine = SyftboxManager.from_config(config.sync, transport)
         # The job client selects a job protocol version per peer from the
         # schema the peer advertises, as the dataset manager does for layouts.
         job_client = JobClient.from_config(

packages/syft-rds/tests/test_dataset_upload_private.py

@@ -8,6 +8,7 @@
     PRIVATE_DATASET_COLLECTION_PREFIX,
 )
 from dataset_test_utils import create_tmp_dataset_files
+from syft.sync.connections.drive.drive_connections import gdrive_connection_of
 
 
 class TestDatasetUploadPrivate:
@@ -145,8 +146,8 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self):
         )
 
         # Get the DS's raw GDrive connection and search for private folders
-        ds_connection: GDriveConnection = (
-            ds_manager.sync_engine._connection_router.connections[0]
+        ds_connection: GDriveConnection = gdrive_connection_of(
+            ds_manager.sync_engine._connection_router.transport
         )
         results = (
             ds_connection.drive_service.files()
@@ -159,8 +160,8 @@ def test_ds_cannot_find_private_folders_via_gdrive_query(self):
         assert len(results.get("files", [])) == 0
 
         # DO should find it with the same query
-        do_connection: GDriveConnection = (
-            do_manager.sync_engine._connection_router.connections[0]
+        do_connection: GDriveConnection = gdrive_connection_of(
+            do_manager.sync_engine._connection_router.transport
         )
         do_results = (
             do_connection.drive_service.files()

packages/syft-rds/tests/test_sync_manager.py

@@ -9,6 +9,7 @@
 
 from syft.sync.connections.drive import mock_drive_service
 from syft.sync.connections.drive.gdrive_transport import GDriveConnection
+from syft.sync.syftbox_manager import SyftboxManager
 from syft_datasets import Dataset
 from syft_datasets.dataset_manager import DATASET_COLLECTION_PREFIX
 from syft_rds import SyftRDSClient
@@ -21,6 +22,7 @@
     create_tmp_dataset_files,
     create_tmp_dataset_files_with_parquet,
 )
+from syft.sync.connections.drive.drive_connections import gdrive_connection_of
 
 
 def path_for_job(do_email: str, ds_email: str, filename: str = "test.job") -> str:
@@ -231,9 +233,9 @@ def test_dataset_only_mock_data_uploaded():
     # Sync so DS receives the dataset
     ds_manager.sync()
 
-    files = ds_manager.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store.files
+    files = gdrive_connection_of(
+        ds_manager.sync_engine._connection_router.transport
+    ).drive_service._backing_store.files
     list(files.keys())
     file_objs = list(files.values())
     file_objs_ex_dataset_yaml = [
@@ -282,7 +284,9 @@ def test_jobs():
     # we do this by not always syncing on a file change, currently this logic is a bit of
     # a short cut, but we could do this based on timing eventually (if there are items in the
     # queue for longer than a certain time we start pushing)
-    connection_do = do_manager.sync_engine._connection_router.connections[0]
+    connection_do = gdrive_connection_of(
+        do_manager.sync_engine._connection_router.transport
+    )
     inbox_folder_id = connection_do._get_own_datasite_inbox_id(ds_manager.email)
     inbox_file_metadatas = connection_do.get_file_metadatas_from_folder(inbox_folder_id)
     assert len(inbox_file_metadatas) == 1
@@ -1036,9 +1040,9 @@ def test_ds_dataset_cache_aware_sync():
     assert remote_hash is not None
 
     # Get the mock backing store and directories for creating second pair
-    mock_backing_store = ds_manager.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store
+    mock_backing_store = gdrive_connection_of(
+        ds_manager.sync_engine._connection_router.transport
+    ).drive_service._backing_store
     ds_folder = ds_manager.syftbox_folder
     do_folder = do_manager.syftbox_folder
     ds_email = ds_manager.email
@@ -1055,12 +1059,12 @@ def test_ds_dataset_cache_aware_sync():
     )
 
     # Replace mock backing store to share dataset collections
-    ds_manager2.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store = mock_backing_store
-    do_manager2.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store = mock_backing_store
+    gdrive_connection_of(
+        ds_manager2.sync_engine._connection_router.transport
+    ).drive_service._backing_store = mock_backing_store
+    gdrive_connection_of(
+        do_manager2.sync_engine._connection_router.transport
+    ).drive_service._backing_store = mock_backing_store
 
     # Load peers (already approved in shared backing store)
     ds_manager2.load_peers()
@@ -1136,9 +1140,9 @@ def test_do_dataset_cache_aware_sync():
     assert len(do_manager.datasets.get_all()) == 1
 
     # Get the mock backing store and directories for creating second pair
-    mock_backing_store = ds_manager.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store
+    mock_backing_store = gdrive_connection_of(
+        ds_manager.sync_engine._connection_router.transport
+    ).drive_service._backing_store
     ds_folder = ds_manager.syftbox_folder
     do_folder = do_manager.syftbox_folder
     ds_email = ds_manager.email
@@ -1155,12 +1159,12 @@ def test_do_dataset_cache_aware_sync():
     )
 
     # Replace mock backing store to share dataset collections
-    ds_manager2.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store = mock_backing_store
-    do_manager2.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store = mock_backing_store
+    gdrive_connection_of(
+        ds_manager2.sync_engine._connection_router.transport
+    ).drive_service._backing_store = mock_backing_store
+    gdrive_connection_of(
+        do_manager2.sync_engine._connection_router.transport
+    ).drive_service._backing_store = mock_backing_store
 
     # Load peers (already approved in shared backing store)
     ds_manager2.load_peers()
@@ -1203,15 +1207,13 @@ def test_in_memory_connection_load_state():
     )
 
     # Get the backing store that will persist across "restarts"
-    backing_store = do_manager1.sync_engine._connection_router.connections[
-        0
-    ].drive_service._backing_store
+    backing_store = gdrive_connection_of(
+        do_manager1.sync_engine._connection_router.transport
+    ).drive_service._backing_store
     ds_folder = ds_manager1.syftbox_folder
     do_folder = do_manager1.syftbox_folder
     ds_email = ds_manager1.email
     do_email = do_manager1.email
-    ds_manager1.sync_engine._connection_router.connections[0]
-    do_manager1.sync_engine._connection_router.connections[0]
 
     # Make some changes (submit to DS's job folder where DS has write access)
     ds_manager1.sync_engine._send_file_change(
@@ -1265,18 +1267,19 @@ def test_in_memory_connection_load_state():
         check_versions=False,
     )
 
-    do_manager2 = SyftRDSClient.from_config(do_config2)
-    ds_manager2 = SyftRDSClient.from_config(ds_config2)
-
     # Connect to the same backing store
     service_do = mock_drive_service.MockDriveService(backing_store, do_email)
     do_connection2 = GDriveConnection.from_service(do_email, service_do)
 
     service_ds = mock_drive_service.MockDriveService(backing_store, ds_email)
     ds_connection2 = GDriveConnection.from_service(ds_email, service_ds)
 
-    do_manager2.sync_engine._add_connection(do_connection2)
-    ds_manager2.sync_engine._add_connection(ds_connection2)
+    do_manager2 = SyftRDSClient.from_config(
+        do_config2, SyftboxManager._drive_transport(do_connection2)
+    )
+    ds_manager2 = SyftRDSClient.from_config(
+        ds_config2, SyftboxManager._drive_transport(ds_connection2)
+    )
 
     # Load peers
     do_manager2.load_peers()

packages/syft-rds/tests/test_version_mismatch_flow.py

@@ -20,6 +20,8 @@
 from syft.version import SYFT_VERSION
 
 from dataset_test_utils import create_test_project_folder, create_tmp_dataset_files
+from syft.sync.connections.drive.drive_connections import gdrive_connection_of
+from syft.sync.syftbox_manager import SyftboxManager
 
 NEW_VERSION = "99.0.0"
 
@@ -46,7 +48,7 @@ def _find_versioned_p2p_folders(connection, peer_email, version):
 
 def _get_backing_store(manager):
     """Extract the mock backing store from a manager's connection."""
-    conn = manager.peer_manager.connection_router.connections[0]
+    conn = gdrive_connection_of(manager.peer_manager.connection_router.transport)
     return conn.drive_service._backing_store
 
 
@@ -66,11 +68,9 @@ def _reinitialize_manager(
         use_in_memory_cache=False,
         syftbox_folder=syftbox_folder,
     )
-    manager = SyftRDSClient.from_config(config)
-
     mock_service = MockDriveService(backing_store, email)
     conn = GDriveConnection.from_service(email, mock_service)
-    manager.sync_engine._add_connection(conn)
+    manager = SyftRDSClient.from_config(config, SyftboxManager._drive_transport(conn))
 
     if has_ds_role:
         manager.sync_engine.file_writer.add_callback(
@@ -172,7 +172,7 @@ def test_version_mismatch_continues_and_repairs():
     assert do_manager.jobs[0].status == "done"
 
     # -- Step 4: Record P2P folders and personal folder id before upgrade --
-    do_conn = do_manager.peer_manager.connection_router.connections[0]
+    do_conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport)
     do_p2p_current = _find_versioned_p2p_folders(
         do_conn, ds_manager.email, SYFT_VERSION
     )
@@ -211,7 +211,9 @@ def test_version_mismatch_continues_and_repairs():
         )
 
         # Personal folder is adopted (same Drive id, new name), not recreated.
-        do_conn_new = do_manager.peer_manager.connection_router.connections[0]
+        do_conn_new = gdrive_connection_of(
+            do_manager.peer_manager.connection_router.transport
+        )
         new_personal_name = f"{NEW_VERSION}#{do_email}"
         new_personal_id = do_conn_new._find_folder_by_name(
             new_personal_name,
@@ -247,7 +249,9 @@ def test_version_mismatch_continues_and_repairs():
             syftbox_folder=ds_syftbox,
         )
 
-        ds_conn_new = ds_manager.peer_manager.connection_router.connections[0]
+        ds_conn_new = gdrive_connection_of(
+            ds_manager.peer_manager.connection_router.transport
+        )
         ds_personal_name = f"{NEW_VERSION}#{ds_email}"
         ds_personal_id = ds_conn_new._find_folder_by_name(
             ds_personal_name,
@@ -327,7 +331,7 @@ def test_login_writes_the_remote_version_file_too():
         use_in_memory_cache=False,
         check_versions=True,
     )
-    conn = do_manager.peer_manager.connection_router.connections[0]
+    conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport)
     assert (
         conn.read_own_peer_info_file().version_info.syft_client_version == SYFT_VERSION
     )
@@ -337,7 +341,9 @@ def test_login_writes_the_remote_version_file_too():
         patch("syft.sync.version.version_info.SYFT_VERSION", NEW_VERSION),
     ):
         upgraded = _upgraded_manager(do_manager)
-        new_conn = upgraded.peer_manager.connection_router.connections[0]
+        new_conn = gdrive_connection_of(
+            upgraded.peer_manager.connection_router.transport
+        )
         # Still the pre-upgrade version: nothing has refreshed it yet.
         assert (
             new_conn.read_own_peer_info_file().version_info.syft_client_version
@@ -388,7 +394,7 @@ def test_the_mismatch_prompt_does_not_return_after_a_login():
             handle_potential_version_mismatches_on_login,
         )
 
-        conn = do_manager.peer_manager.connection_router.connections[0]
+        conn = gdrive_connection_of(do_manager.peer_manager.connection_router.transport)
         with patch(
             "syft.sync.login_utils._read_remote_version",
             side_effect=lambda e, t: conn.read_own_peer_info_file().version_info,