From 9b2bf690fd231687dde8ecfa247b7443e98c8b4b Mon Sep 17 00:00:00 2001
From: Aaron Taylor <aaron@additive.ai>
Date: Tue, 3 Sep 2024 18:54:36 -0700
Subject: [PATCH] update custom scheme tests to utilize pytest fixture

This isolates the implementation in response to PR feedback:
https://github.com/drivendataorg/cloudpathlib/pull/467#discussion_r1742723302
---
 tests/test_client.py | 65 ++++++++++++++++++++++++++++++--------------
 1 file changed, 45 insertions(+), 20 deletions(-)

diff --git a/tests/test_client.py b/tests/test_client.py
index cb5b5a94..af0206a8 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -4,15 +4,19 @@
 import string
 from pathlib import Path
 
+import pytest
+
 from cloudpathlib import CloudPath
 from cloudpathlib.client import register_client_class
-from cloudpathlib.cloudpath import register_path_class
+from cloudpathlib.cloudpath import implementation_registry, register_path_class
 from cloudpathlib.s3.s3client import S3Client
 from cloudpathlib.s3.s3path import S3Path
 
 
 def test_default_client_instantiation(rig):
-    if not getattr(rig, "is_custom_s3", False) and not (getattr(rig, "is_adls_gen2", False)):
+    if not getattr(rig, "is_custom_s3", False) and not (
+        getattr(rig, "is_adls_gen2", False)
+    ):
         # Skip resetting the default client for custom S3 endpoint, but keep the other tests,
         # since they're still useful.
         rig.client_class._default_client = None
@@ -47,7 +51,9 @@ def test_different_clients(rig):
     p = rig.create_cloud_path("dir_0/file0_0.txt")
 
     new_client = rig.client_class(**rig.required_client_kwargs)
-    p2 = new_client.CloudPath(f"{rig.cloud_prefix}{rig.drive}/{rig.test_dir}/dir_0/file0_0.txt")
+    p2 = new_client.CloudPath(
+        f"{rig.cloud_prefix}{rig.drive}/{rig.test_dir}/dir_0/file0_0.txt"
+    )
 
     assert p.client is not p2.client
     assert p._local is not p2._local
@@ -106,7 +112,9 @@ def my_content_type(path):
 
     # see if testing custom s3 endpoint, make sure to pass the url to the constructor
     kwargs = rig.required_client_kwargs.copy()
-    custom_endpoint = os.getenv("CUSTOM_S3_ENDPOINT", "https://s3.us-west-1.drivendatabws.com")
+    custom_endpoint = os.getenv(
+        "CUSTOM_S3_ENDPOINT", "https://s3.us-west-1.drivendatabws.com"
+    )
     if (
         rig.client_class is S3Client
         and rig.live_server
@@ -115,38 +123,55 @@ def my_content_type(path):
         kwargs["endpoint_url"] = custom_endpoint
 
     # set up default client to use content_type_method
-    rig.client_class(content_type_method=my_content_type, **kwargs).set_as_default_client()
+    rig.client_class(
+        content_type_method=my_content_type, **kwargs
+    ).set_as_default_client()
 
     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://"
+@pytest.fixture
+def custom_s3_path():
+    # A fixture isolates these classes as they modify the global registry of
+    # implementations.
+    @register_path_class("mys3")
+    class MyS3Path(S3Path):
+        cloud_prefix: str = "mys3://"
 
+    @register_client_class("mys3")
+    class MyS3Client(S3Client):
+        pass
 
-@register_client_class("mys3")
-class MyS3Client(S3Client):
-    pass
+    yield (MyS3Path, MyS3Client)
 
+    # cleanup after use
+    implementation_registry.pop("mys3")
 
-def test_custom_mys3path_instantiation():
-    path = MyS3Path("mys3://bucket/dir/file.txt")
-    assert isinstance(path, MyS3Path)
+
+def test_custom_mys3path_instantiation(custom_s3_path):
+    CustomPath, _ = custom_s3_path
+
+    path = CustomPath("mys3://bucket/dir/file.txt")
+    assert isinstance(path, CustomPath)
     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)
+def test_custom_mys3client_instantiation(custom_s3_path):
+    _, CustomClient = custom_s3_path
+
+    client = CustomClient()
+    assert isinstance(client, CustomClient)
     assert client.CloudPath("mys3://bucket/dir/file.txt").cloud_prefix == "mys3://"
 
 
-def test_custom_mys3client_default_client():
-    MyS3Client().set_as_default_client()
+def test_custom_mys3client_default_client(custom_s3_path):
+    _, CustomClient = custom_s3_path
+
+    CustomClient().set_as_default_client()
+
     path = CloudPath("mys3://bucket/dir/file.txt")
-    assert isinstance(path.client, MyS3Client)
+    assert isinstance(path.client, CustomClient)
     assert path.cloud_prefix == "mys3://"