From 672f724a282aa4e9f3ab44fa6d9d9666f1bdb275 Mon Sep 17 00:00:00 2001
From: Aaron Taylor <aaron@additive.ai>
Date: Tue, 3 Sep 2024 13:40:56 -0700
Subject: [PATCH] tests: switch to lighter-weight custom scheme tests

---
 tests/conftest.py    | 45 +++-----------------------------------------
 tests/test_client.py | 36 ++++++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 43 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 6b742be8..a83046bf 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -27,8 +27,7 @@
     S3Path,
 )
 from cloudpathlib.azure.azblobclient import _hns_rmtree
-from cloudpathlib.client import register_client_class
-from cloudpathlib.cloudpath import implementation_registry, register_path_class
+from cloudpathlib.cloudpath import implementation_registry
 from cloudpathlib.local import (
     LocalAzureBlobClient,
     LocalAzureBlobPath,
@@ -317,7 +316,8 @@ def s3_rig(request, monkeypatch, assets_dir):
         bucket.objects.filter(Prefix=test_dir).delete()
 
 
-def _custom_s3_rig_helper(request, monkeypatch, assets_dir, path_class, client_class):
+@fixture()
+def custom_s3_rig(request, monkeypatch, assets_dir):
     """
     Custom S3 rig used to test the integrations with non-AWS S3-compatible object storages like
         - MinIO (https://min.io/)
@@ -402,43 +402,6 @@ def _spin_up_bucket():
         bucket.objects.filter(Prefix=test_dir).delete()
 
 
-@fixture()
-def custom_s3_rig(request, monkeypatch, assets_dir):
-    """
-    Custom S3 rig used to test the integrations with non-AWS S3-compatible object storages like
-        - MinIO (https://min.io/)
-        - CEPH  (https://ceph.io/ceph-storage/object-storage/)
-        - others
-    """
-    yield from _custom_s3_rig_helper(request, monkeypatch, assets_dir, S3Path, S3Client)
-
-
-@register_path_class("mys3")
-class MyS3Path(S3Path):
-    cloud_prefix: str = "mys3://"
-
-
-@register_client_class("mys3")
-class MyS3Client(S3Client):
-    pass
-
-
-# Mirrors the definition of the S3Client class
-MyS3Client.MyS3Path = MyS3Client.CloudPath  # type: ignore
-
-
-@fixture()
-def custom_scheme_s3_rig(request, monkeypatch, assets_dir):
-    """
-    Custom S3 rig used to test the integrations with non-AWS S3-compatible object storages like
-        - MinIO (https://min.io/)
-        - CEPH  (https://ceph.io/ceph-storage/object-storage/)
-        - others
-    with the addition of a custom scheme being used.
-    """
-    yield from _custom_s3_rig_helper(request, monkeypatch, assets_dir, MyS3Path, MyS3Client)
-
-
 @fixture()
 def local_azure_rig(request, monkeypatch, assets_dir):
     drive = os.getenv("LIVE_AZURE_CONTAINER", DEFAULT_CONTAINER_NAME)
@@ -532,7 +495,6 @@ def local_s3_rig(request, monkeypatch, assets_dir):
         gs_rig,
         s3_rig,
         custom_s3_rig,
-        custom_scheme_s3_rig,
         local_azure_rig,
         local_s3_rig,
         local_gs_rig,
@@ -545,6 +507,5 @@ def local_s3_rig(request, monkeypatch, assets_dir):
     [
         s3_rig,
         custom_s3_rig,
-        custom_scheme_s3_rig,
     ],
 )
diff --git a/tests/test_client.py b/tests/test_client.py
index a665a5a6..cb5b5a94 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -1,11 +1,14 @@
 import mimetypes
 import os
-from pathlib import Path
 import random
 import string
+from pathlib import Path
 
 from cloudpathlib import CloudPath
+from cloudpathlib.client import register_client_class
+from cloudpathlib.cloudpath import register_path_class
 from cloudpathlib.s3.s3client import S3Client
+from cloudpathlib.s3.s3path import S3Path
 
 
 def test_default_client_instantiation(rig):
@@ -116,3 +119,34 @@ def my_content_type(path):
 
     for suffix, content_type in mimes:
         _test_write_content_type(suffix, content_type, rig)
+
+
+@register_path_class("mys3")
+class MyS3Path(S3Path):
+    cloud_prefix: str = "mys3://"
+
+
+@register_client_class("mys3")
+class MyS3Client(S3Client):
+    pass
+
+
+def test_custom_mys3path_instantiation():
+    path = MyS3Path("mys3://bucket/dir/file.txt")
+    assert isinstance(path, MyS3Path)
+    assert path.cloud_prefix == "mys3://"
+    assert path.bucket == "bucket"
+    assert path.key == "dir/file.txt"
+
+
+def test_custom_mys3client_instantiation():
+    client = MyS3Client()
+    assert isinstance(client, MyS3Client)
+    assert client.CloudPath("mys3://bucket/dir/file.txt").cloud_prefix == "mys3://"
+
+
+def test_custom_mys3client_default_client():
+    MyS3Client().set_as_default_client()
+    path = CloudPath("mys3://bucket/dir/file.txt")
+    assert isinstance(path.client, MyS3Client)
+    assert path.cloud_prefix == "mys3://"