Increase coverage of existing functions

Author: kylstevenson

py_tests/test_authentication.py

@@ -9,13 +9,15 @@
 
 import fsspec
 import pytest
+import requests
 from requests.adapters import HTTPAdapter
 
 from fusion._legacy.authentication import (
     try_get_client_id,
     try_get_client_secret,
 )
 from fusion.authentication import (
+    FusionAiohttpSession,
     FusionOAuthAdapter,
 )
 from fusion.credentials import FusionCredentials
@@ -264,3 +266,70 @@ def test_send_raises_api_response_error_on_connection_error() -> None:
     assert "Connection error while sending request" in str(exc_info.value)
     HTTP_STATUS_SERVICE_UNAVAILABLE = 503
     assert exc_info.value.status_code == HTTP_STATUS_SERVICE_UNAVAILABLE
+
+
+def test_fusion_oauth_adapter_uses_credentials_proxies_and_default_retry(credentials: FusionCredentials) -> None:
+    default_retry_total = 20
+    adapter = FusionOAuthAdapter(credentials=credentials)
+
+    assert adapter.proxies == credentials.proxies
+    assert adapter.auth_retries.total == default_retry_total
+    assert adapter.headers == {}
+
+
+def test_fusion_oauth_adapter_send_merges_auth_and_custom_headers() -> None:
+    request = Mock()
+    request.url = "https://example.com/catalogs/c1/datasets/d1"
+    request.headers = {}
+
+    mock_credentials = Mock(spec=FusionCredentials)
+    mock_credentials.proxies = {}
+    mock_credentials.get_fusion_token_headers.return_value = {"Authorization": "Bearer abc"}
+
+    adapter = FusionOAuthAdapter(credentials=mock_credentials, headers={"X-Test": "1"})
+
+    response = requests.Response()
+    response.status_code = 200
+    response._content = b"ok"
+
+    with patch.object(HTTPAdapter, "send", return_value=response) as mock_send:
+        result = adapter.send(request)
+
+    assert result is response
+    assert request.headers["Authorization"] == "Bearer abc"
+    assert request.headers["X-Test"] == "1"
+    mock_send.assert_called_once()
+
+
+def test_send_raises_api_response_error_when_header_generation_fails() -> None:
+    request = Mock()
+    request.url = "https://example.com/data"
+    request.headers = {}
+
+    mock_credentials = Mock(spec=FusionCredentials)
+    mock_credentials.proxies = {}
+    mock_credentials.get_fusion_token_headers.side_effect = RuntimeError("bad headers")
+
+    adapter = FusionOAuthAdapter(credentials=mock_credentials)
+
+    with pytest.raises(APIResponseError, match="Failed to generate Fusion token headers") as exc_info:
+        adapter.send(request)
+
+    internal_server_error = 500
+    assert exc_info.value.status_code == internal_server_error
+
+
+@pytest.mark.asyncio
+async def test_fusion_aiohttp_session_post_init_sets_state(credentials: FusionCredentials) -> None:
+    refresh_within_seconds = 9
+    session = FusionAiohttpSession()
+    try:
+        session.post_init(credentials=credentials, refresh_within_seconds=refresh_within_seconds)
+
+        assert session.credentials is credentials
+        assert session.refresh_within_seconds == refresh_within_seconds
+        assert session.number_token_refreshes == 0
+        assert session.fusion_token_dict == {}
+        assert session.fusion_token_expiry_dict == {}
+    finally:
+        await session.close()

py_tests/test_cli_and_shims.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from pathlib import Path
+from queue import Queue
+from typing import TYPE_CHECKING, Any, get_args
+
+if TYPE_CHECKING:
+    import pytest
+
+
+class TestFusionMainModule:
+    @staticmethod
+    def _run_main_module() -> None:
+        main_path = Path(__file__).resolve().parents[1] / "py_src" / "fusion" / "__main__.py"
+        namespace = {"__file__": str(main_path), "__name__": "__main__"}
+        exec(compile(main_path.read_text(), str(main_path), "exec"), namespace)  # noqa: S102
+
+    def test_invokes_requested_method(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        import fusion
+
+        captured: dict[str, Any] = {}
+
+        class FakeFusion:
+            def __init__(
+                self,
+                root_url: str | None = None,
+                credentials: str | None = None,
+                download_folder: str | None = None,
+                log_level: str | None = None,
+                log_path: str | None = None,
+            ) -> None:
+                captured["init"] = {
+                    "root_url": root_url,
+                    "credentials": credentials,
+                    "download_folder": download_folder,
+                    "log_level": log_level,
+                    "log_path": log_path,
+                }
+
+            def publish(self, dataset: str | None = None, overwrite: bool = False, dry_run: bool = True) -> None:
+                captured["method"] = {
+                    "dataset": dataset,
+                    "overwrite": overwrite,
+                    "dry_run": dry_run,
+                }
+
+        monkeypatch.setattr(fusion, "Fusion", FakeFusion)
+        monkeypatch.setattr(
+            "sys.argv",
+            [
+                "fusion",
+                "--root_url",
+                "https://example.test/api/v1/",
+                "--credentials",
+                "client_credentials.json",
+                "--download_folder",
+                "downloads-dir",
+                "--log_level",
+                "20",
+                "--log_path",
+                "logs",
+                "--method",
+                "publish",
+                "--dataset",
+                "prices",
+                "--overwrite",
+                "True",
+                "--dry_run",
+                "False",
+            ],
+        )
+
+        self._run_main_module()
+
+        assert captured["init"] == {
+            "root_url": "https://example.test/api/v1/",
+            "credentials": "client_credentials.json",
+            "download_folder": "downloads-dir",
+            "log_level": "20",
+            "log_path": "logs",
+        }
+        assert captured["method"] == {
+            "dataset": "prices",
+            "overwrite": True,
+            "dry_run": False,
+        }
+
+    def test_without_method_only_initializes_client(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        import fusion
+
+        captured: dict[str, Any] = {}
+
+        class FakeFusion:
+            def __init__(
+                self,
+                root_url: str | None = None,
+                credentials: str | None = None,
+                download_folder: str | None = None,
+                log_level: str | None = None,
+                log_path: str | None = None,
+            ) -> None:
+                captured["init"] = {
+                    "root_url": root_url,
+                    "credentials": credentials,
+                    "download_folder": download_folder,
+                    "log_level": log_level,
+                    "log_path": log_path,
+                }
+
+            def publish(self, dataset: str | None = None) -> None:
+                captured["method"] = dataset
+
+        monkeypatch.setattr(fusion, "Fusion", FakeFusion)
+        monkeypatch.setattr("sys.argv", ["fusion", "--credentials", "client_credentials.json"])
+
+        self._run_main_module()
+
+        assert captured["init"] == {
+            "root_url": None,
+            "credentials": "client_credentials.json",
+            "download_folder": None,
+            "log_level": None,
+            "log_path": None,
+        }
+        assert "method" not in captured
+
+
+class TestCompatibilityShims:
+    def test_legacy_fusion_module_exports_fusion_credentials(self) -> None:
+        shim_path = Path(__file__).resolve().parents[1] / "py_src" / "fusion" / "_fusion.py"
+        namespace: dict[str, Any] = {"__file__": str(shim_path), "__name__": "fusion._fusion_py"}
+
+        exec(compile(shim_path.read_text(), str(shim_path), "exec"), namespace)  # noqa: S102
+
+        assert namespace["FusionCredentials"].__name__ == "FusionCredentials"
+        assert namespace["__all__"] == ["FusionCredentials"]
+
+    def test_types_aliases_are_available(self) -> None:
+        from fusion.types import PyArrowFilterT, WorkerQueueT
+
+        assert WorkerQueueT == Queue[tuple[int, int, int]]
+        assert get_args(PyArrowFilterT) == (list[tuple[Any]], list[list[tuple[Any]]])