add test and refactor

Author: Coniferish

_test_unstructured_client/test_check_url_protocol.py renamed to _test_unstructured_client/test__decorators.py

@@ -1,15 +1,11 @@
-import os
 import pytest
 
 from unstructured_client import UnstructuredClient
+from unstructured_client.models import shared
+from unstructured_client.models.errors import SDKError
 
 
-def get_api_key():
-    api_key = os.getenv("UNS_API_KEY")
-    if api_key is None:
-        raise ValueError("""UNS_API_KEY environment variable not set. 
-Set it in your current shell session with `export UNS_API_KEY=<api_key>`""")
-    return api_key
+FAKE_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
 
 
 @pytest.mark.parametrize(
@@ -25,7 +21,7 @@ def get_api_key():
 def test_clean_server_url_on_paid_api_url(server_url: str):
     client = UnstructuredClient(
         server_url=server_url,
-        api_key_auth=get_api_key(),
+        api_key_auth=FAKE_KEY,
     )
     assert client.general.sdk_configuration.server_url == "https://unstructured-000mock.api.unstructuredapp.io"
 
@@ -42,18 +38,19 @@ def test_clean_server_url_on_paid_api_url(server_url: str):
 def test_clean_server_url_on_localhost(server_url: str):
     client = UnstructuredClient(
         server_url=server_url,
-        api_key_auth=get_api_key(),
+        api_key_auth=FAKE_KEY,
     )
     assert client.general.sdk_configuration.server_url == "http://localhost:8000"
 
 
 def test_clean_server_url_on_empty_string():
     client = UnstructuredClient(
         server_url="",
-        api_key_auth=get_api_key(),
+        api_key_auth=FAKE_KEY,
     )
     assert client.general.sdk_configuration.server_url == ""
 
+
 @pytest.mark.parametrize(
     ("server_url"),
     [
@@ -63,8 +60,33 @@ def test_clean_server_url_on_empty_string():
 )
 def test_clean_server_url_with_positional_arguments(server_url: str):
     client = UnstructuredClient(
-        get_api_key(),
+        FAKE_KEY,
         "",
         server_url,
     )
     assert client.general.sdk_configuration.server_url == "https://unstructured-000mock.api.unstructuredapp.io"
+
+
+def test_suggest_defining_url_if_401():
+    with pytest.warns(UserWarning):
+
+        client = UnstructuredClient(
+            api_key_auth=FAKE_KEY,
+        )
+        
+        filename = "_sample_docs/layout-parser-paper-fast.pdf"
+        
+        with open(filename, "rb") as f:
+            files=shared.Files(
+                content=f.read(),
+                file_name=filename,
+            )
+
+        req = shared.PartitionParameters(
+            files=files,
+        )
+
+        try:
+            client.general.partition(req)
+        except SDKError as e:
+            print(e)

src/unstructured_client/utils/_decorators.py

@@ -6,7 +6,7 @@
 from urllib.parse import urlparse, urlunparse, ParseResult
 import warnings
 
-from unstructured_client.models import errors
+from unstructured_client.models import errors, operations
 
 if TYPE_CHECKING:
     from unstructured_client.general import General
@@ -52,14 +52,14 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
     return wrapper
 
 
-def suggest_defining_url_if_401(func: Callable[_P, None]) -> Callable[_P, None]:
+def suggest_defining_url_if_401(func: Callable[_P, operations.PartitionResponse]) -> Callable[_P, operations.PartitionResponse]:
 
     @functools.wraps(func)
-    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> operations.PartitionResponse:
         try:
             return func(*args, **kwargs)
         except errors.SDKError:
-            general_obj: General = args[0]
+            general_obj: General = args[0] # type: ignore
             if not general_obj.sdk_configuration.server_url:
                 warnings.warn("If intending to use the paid API, please define `server_url` in your request.")
             return func(*args, **kwargs)