Use shared upload file context helper in session managers

Co-authored-by: Shri Sukhani <shrisukhani@users.noreply.github.com>

Author: cursoragent

hyperbrowser/client/managers/async_manager/session.py

@@ -1,15 +1,13 @@
 from os import PathLike
 from typing import IO, List, Optional, Union, overload
 import warnings
-from hyperbrowser.exceptions import HyperbrowserError
-from ...file_utils import open_binary_file
 from ..serialization_utils import (
     serialize_model_dump_or_default,
     serialize_model_dump_to_dict,
     serialize_optional_model_dump_to_dict,
 )
 from ..session_profile_update_utils import resolve_update_profile_params
-from ..session_upload_utils import normalize_upload_file_input
+from ..session_upload_utils import open_upload_files_from_input
 from ..session_utils import (
     parse_session_recordings_response_data,
     parse_session_response_model,
@@ -168,23 +166,7 @@ async def get_downloads_url(self, id: str) -> GetSessionDownloadsUrlResponse:
     async def upload_file(
         self, id: str, file_input: Union[str, PathLike[str], IO]
     ) -> UploadFileResponse:
-        file_path, file_obj = normalize_upload_file_input(file_input)
-        if file_path is not None:
-            with open_binary_file(
-                file_path,
-                open_error_message=f"Failed to open upload file at path: {file_path}",
-            ) as file_obj:
-                files = {"file": file_obj}
-                response = await self._client.transport.post(
-                    self._client._build_url(f"/session/{id}/uploads"),
-                    files=files,
-                )
-        else:
-            if file_obj is None:
-                raise HyperbrowserError(
-                    "file_input must be a file path or file-like object"
-                )
-            files = {"file": file_obj}
+        with open_upload_files_from_input(file_input) as files:
             response = await self._client.transport.post(
                 self._client._build_url(f"/session/{id}/uploads"),
                 files=files,

hyperbrowser/client/managers/session_upload_utils.py

@@ -1,11 +1,12 @@
 import os
+from contextlib import contextmanager
 from os import PathLike
-from typing import IO, Optional, Tuple, Union
+from typing import Dict, IO, Iterator, Optional, Tuple, Union
 
 from hyperbrowser.exceptions import HyperbrowserError
 from hyperbrowser.type_utils import is_plain_string, is_string_subclass_instance
 
-from ..file_utils import ensure_existing_file_path
+from ..file_utils import ensure_existing_file_path, open_binary_file
 
 
 def normalize_upload_file_input(
@@ -56,3 +57,20 @@ def normalize_upload_file_input(
         raise HyperbrowserError("file_input file-like object must be open")
 
     return None, file_input
+
+
+@contextmanager
+def open_upload_files_from_input(
+    file_input: Union[str, PathLike[str], IO],
+) -> Iterator[Dict[str, IO]]:
+    file_path, file_obj = normalize_upload_file_input(file_input)
+    if file_path is not None:
+        with open_binary_file(
+            file_path,
+            open_error_message=f"Failed to open upload file at path: {file_path}",
+        ) as opened_file:
+            yield {"file": opened_file}
+        return
+    if file_obj is None:
+        raise HyperbrowserError("file_input must be a file path or file-like object")
+    yield {"file": file_obj}

hyperbrowser/client/managers/sync_manager/session.py

@@ -1,15 +1,13 @@
 from os import PathLike
 from typing import IO, List, Optional, Union, overload
 import warnings
-from hyperbrowser.exceptions import HyperbrowserError
-from ...file_utils import open_binary_file
 from ..serialization_utils import (
     serialize_model_dump_or_default,
     serialize_model_dump_to_dict,
     serialize_optional_model_dump_to_dict,
 )
 from ..session_profile_update_utils import resolve_update_profile_params
-from ..session_upload_utils import normalize_upload_file_input
+from ..session_upload_utils import open_upload_files_from_input
 from ..session_utils import (
     parse_session_recordings_response_data,
     parse_session_response_model,
@@ -160,23 +158,7 @@ def get_downloads_url(self, id: str) -> GetSessionDownloadsUrlResponse:
     def upload_file(
         self, id: str, file_input: Union[str, PathLike[str], IO]
     ) -> UploadFileResponse:
-        file_path, file_obj = normalize_upload_file_input(file_input)
-        if file_path is not None:
-            with open_binary_file(
-                file_path,
-                open_error_message=f"Failed to open upload file at path: {file_path}",
-            ) as file_obj:
-                files = {"file": file_obj}
-                response = self._client.transport.post(
-                    self._client._build_url(f"/session/{id}/uploads"),
-                    files=files,
-                )
-        else:
-            if file_obj is None:
-                raise HyperbrowserError(
-                    "file_input must be a file path or file-like object"
-                )
-            files = {"file": file_obj}
+        with open_upload_files_from_input(file_input) as files:
             response = self._client.transport.post(
                 self._client._build_url(f"/session/{id}/uploads"),
                 files=files,

tests/test_binary_file_open_helper_usage.py

@@ -5,16 +5,27 @@
 pytestmark = pytest.mark.architecture
 
 
-MODULES = (
+EXTENSION_MODULES = (
     "hyperbrowser/client/managers/sync_manager/extension.py",
     "hyperbrowser/client/managers/async_manager/extension.py",
+)
+
+SESSION_MODULES = (
     "hyperbrowser/client/managers/sync_manager/session.py",
     "hyperbrowser/client/managers/async_manager/session.py",
 )
 
 
-def test_managers_use_shared_binary_file_open_helper():
-    for module_path in MODULES:
+def test_extension_managers_use_shared_binary_file_open_helper():
+    for module_path in EXTENSION_MODULES:
         module_text = Path(module_path).read_text(encoding="utf-8")
         assert "open_binary_file(" in module_text
         assert 'with open(file_path, "rb")' not in module_text
+
+
+def test_session_managers_use_upload_file_context_helper():
+    for module_path in SESSION_MODULES:
+        module_text = Path(module_path).read_text(encoding="utf-8")
+        assert "open_upload_files_from_input(" in module_text
+        assert "open_binary_file(" not in module_text
+        assert 'with open(file_path, "rb")' not in module_text

tests/test_session_upload_helper_usage.py

@@ -14,7 +14,8 @@
 def test_session_managers_use_shared_upload_input_normalizer():
     for module_path in SESSION_MANAGER_MODULES:
         module_text = Path(module_path).read_text(encoding="utf-8")
-        assert "normalize_upload_file_input(" in module_text
+        assert "open_upload_files_from_input(" in module_text
         assert "os.fspath(" not in module_text
         assert "ensure_existing_file_path(" not in module_text
         assert 'getattr(file_input, "read"' not in module_text
+        assert "open_binary_file(" not in module_text

tests/test_session_upload_utils.py

@@ -4,7 +4,9 @@
 
 import pytest
 
+import hyperbrowser.client.managers.session_upload_utils as session_upload_utils
 from hyperbrowser.client.managers.session_upload_utils import (
+    open_upload_files_from_input,
     normalize_upload_file_input,
 )
 from hyperbrowser.exceptions import HyperbrowserError
@@ -134,3 +136,38 @@ def closed(self):
         normalize_upload_file_input(_BrokenFileLike())  # type: ignore[arg-type]
 
     assert exc_info.value.original_error is None
+
+
+def test_open_upload_files_from_input_opens_and_closes_path_input(tmp_path: Path):
+    file_path = tmp_path / "file.txt"
+    file_path.write_text("content")
+
+    with open_upload_files_from_input(str(file_path)) as files:
+        assert "file" in files
+        assert files["file"].closed is False
+        assert files["file"].read() == b"content"
+
+    assert files["file"].closed is True
+
+
+def test_open_upload_files_from_input_reuses_file_like_object():
+    file_obj = io.BytesIO(b"content")
+
+    with open_upload_files_from_input(file_obj) as files:
+        assert files == {"file": file_obj}
+
+
+def test_open_upload_files_from_input_rejects_missing_normalized_file_object(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(
+        session_upload_utils,
+        "normalize_upload_file_input",
+        lambda file_input: (None, None),
+    )
+
+    with pytest.raises(
+        HyperbrowserError, match="file_input must be a file path or file-like object"
+    ):
+        with open_upload_files_from_input(io.BytesIO(b"content")):
+            pass