Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 49 additions & 1 deletion src/runloop_api_client/sdk/async_.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,12 @@

from __future__ import annotations

import io
import asyncio
import tarfile
from typing import Dict, Mapping, Optional
from pathlib import Path
from datetime import timedelta
from typing_extensions import Unpack

import httpx
Expand Down Expand Up @@ -350,7 +354,7 @@ async def upload_from_file(
path = Path(file_path)

try:
content = path.read_bytes()
content = await asyncio.to_thread(lambda: path.read_bytes())
except OSError as error:
raise OSError(f"Failed to read file {path}: {error}") from error

Expand All @@ -361,6 +365,50 @@ async def upload_from_file(
await obj.complete()
return obj

async def upload_from_dir(
self,
dir_path: str | Path,
*,
name: Optional[str] = None,
metadata: Optional[Dict[str, str]] = None,
ttl: Optional[timedelta] = None,
**options: Unpack[LongRequestOptions],
) -> AsyncStorageObject:
"""Create and upload an object from a local directory.

The resulting object will be uploaded as a compressed tarball.

:param dir_path: Local filesystem directory path to tar
:type dir_path: str | Path
:param name: Optional object name; defaults to the directory name + '.tar.gz'
:type name: Optional[str]
:param metadata: Optional key-value metadata
:type metadata: Optional[Dict[str, str]]
:param ttl: Optional Time-To-Live, after which the object is automatically deleted
:type ttl: Optional[timedelta]
:param options: See :typeddict:`~runloop_api_client.sdk._types.LongRequestOptions` for available options
:return: Wrapper for the uploaded object
:rtype: AsyncStorageObject
:raises OSError: If the local file cannot be read
"""
path = Path(dir_path)
name = name or f"{path.name}.tar.gz"
ttl_ms = int(ttl.total_seconds()) * 1000 if ttl else None

def synchronous_io() -> bytes:
with io.BytesIO() as tar_buffer:
with tarfile.open(fileobj=tar_buffer, mode="w:gz") as tar:
tar.add(path, arcname=".", recursive=True)
tar_buffer.seek(0)
return tar_buffer.read()

tar_bytes = await asyncio.to_thread(synchronous_io)

obj = await self.create(name=name, content_type="tgz", metadata=metadata, ttl_ms=ttl_ms, **options)
await obj.upload_content(tar_bytes)
await obj.complete()
return obj

async def upload_from_text(
self,
text: str,
Expand Down
3 changes: 2 additions & 1 deletion src/runloop_api_client/sdk/async_storage_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from __future__ import annotations

from typing import Iterable
from typing_extensions import Unpack, override

from ._types import RequestOptions, LongRequestOptions, SDKObjectDownloadParams
Expand Down Expand Up @@ -146,7 +147,7 @@ async def delete(
**options,
)

async def upload_content(self, content: str | bytes) -> None:
async def upload_content(self, content: str | bytes | Iterable[bytes]) -> None:
"""Upload content to the object's pre-signed URL.

:param content: Bytes or text payload to upload
Expand Down
3 changes: 2 additions & 1 deletion src/runloop_api_client/sdk/storage_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from __future__ import annotations

from typing import Iterable
from typing_extensions import Unpack, override

from ._types import RequestOptions, LongRequestOptions, SDKObjectDownloadParams
Expand Down Expand Up @@ -146,7 +147,7 @@ def delete(
**options,
)

def upload_content(self, content: str | bytes) -> None:
def upload_content(self, content: str | bytes | Iterable[bytes]) -> None:
"""Upload content to the object's pre-signed URL.

:param content: Bytes or text payload to upload
Expand Down
43 changes: 43 additions & 0 deletions src/runloop_api_client/sdk/sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,11 @@

from __future__ import annotations

import io
import tarfile
from typing import Dict, Mapping, Optional
from pathlib import Path
from datetime import timedelta
from typing_extensions import Unpack

import httpx
Expand Down Expand Up @@ -361,6 +364,46 @@ def upload_from_file(
obj.complete()
return obj

def upload_from_dir(
self,
dir_path: str | Path,
*,
name: Optional[str] = None,
metadata: Optional[Dict[str, str]] = None,
ttl: Optional[timedelta] = None,
**options: Unpack[LongRequestOptions],
) -> StorageObject:
"""Create and upload an object from a local directory.

The resulting object will be uploaded as a compressed tarball.

:param dir_path: Local filesystem directory path to tar
:type dir_path: str | Path
:param name: Optional object name; defaults to the directory name + '.tar.gz'
:type name: Optional[str]
:param metadata: Optional key-value metadata
:type metadata: Optional[Dict[str, str]]
:param ttl: Optional Time-To-Live, after which the object is automatically deleted
:type ttl: Optional[timedelta]
:param options: See :typeddict:`~runloop_api_client.sdk._types.LongRequestOptions` for available options
:return: Wrapper for the uploaded object
:rtype: StorageObject
:raises OSError: If the local file cannot be read
"""
path = Path(dir_path)
name = name or f"{path.name}.tar.gz"
ttl_ms = int(ttl.total_seconds()) * 1000 if ttl else None

tar_buffer = io.BytesIO()
with tarfile.open(fileobj=tar_buffer, mode="w:gz") as tar:
tar.add(path, arcname=".", recursive=True)
tar_buffer.seek(0)

obj = self.create(name=name, content_type="tgz", metadata=metadata, ttl_ms=ttl_ms, **options)
obj.upload_content(tar_buffer)
obj.complete()
return obj

def upload_from_text(
self,
text: str,
Expand Down
157 changes: 157 additions & 0 deletions tests/sdk/test_async_clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

from __future__ import annotations

import io
import tarfile
from types import SimpleNamespace
from pathlib import Path
from unittest.mock import AsyncMock
Expand Down Expand Up @@ -327,6 +329,161 @@ async def test_upload_from_file_missing_path(self, mock_async_client: AsyncMock,
with pytest.raises(OSError, match="Failed to read file"):
await client.upload_from_file(missing_file)

@pytest.mark.asyncio
async def test_upload_from_dir(
self, mock_async_client: AsyncMock, object_view: MockObjectView, tmp_path: Path
) -> None:
"""Test upload_from_dir method."""
mock_async_client.objects.create = AsyncMock(return_value=object_view)
mock_async_client.objects.complete = AsyncMock(return_value=object_view)

# Create a temporary directory with some files
test_dir = tmp_path / "test_directory"
test_dir.mkdir()
(test_dir / "file1.txt").write_text("content1")
(test_dir / "file2.txt").write_text("content2")
subdir = test_dir / "subdir"
subdir.mkdir()
(subdir / "file3.txt").write_text("content3")

http_client = AsyncMock()
mock_response = create_mock_httpx_response()
http_client.put = AsyncMock(return_value=mock_response)
mock_async_client._client = http_client

client = AsyncStorageObjectOps(mock_async_client)
obj = await client.upload_from_dir(test_dir, name="archive.tar.gz", metadata={"key": "value"})

assert isinstance(obj, AsyncStorageObject)
assert obj.id == "obj_123"
mock_async_client.objects.create.assert_awaited_once_with(
name="archive.tar.gz",
content_type="tgz",
metadata={"key": "value"},
ttl_ms=None,
)
# Verify that put was called with tarball content
http_client.put.assert_awaited_once()
call_args = http_client.put.call_args
assert call_args[0][0] == object_view.upload_url

# Verify it's a valid gzipped tarball
uploaded_content = call_args[1]["content"]
with tarfile.open(fileobj=io.BytesIO(uploaded_content), mode="r:gz") as tar:
members = tar.getmembers()
member_names = [m.name for m in members]
# Should contain our test files (may include directory entries)
assert any("file1.txt" in name for name in member_names)
assert any("file2.txt" in name for name in member_names)
assert any("file3.txt" in name for name in member_names)

mock_async_client.objects.complete.assert_awaited_once()

@pytest.mark.asyncio
async def test_upload_from_dir_default_name(
self, mock_async_client: AsyncMock, object_view: MockObjectView, tmp_path: Path
) -> None:
"""Test upload_from_dir uses directory name by default."""
mock_async_client.objects.create = AsyncMock(return_value=object_view)
mock_async_client.objects.complete = AsyncMock(return_value=object_view)

test_dir = tmp_path / "my_folder"
test_dir.mkdir()
(test_dir / "file.txt").write_text("content")

http_client = AsyncMock()
mock_response = create_mock_httpx_response()
http_client.put = AsyncMock(return_value=mock_response)
mock_async_client._client = http_client

client = AsyncStorageObjectOps(mock_async_client)
obj = await client.upload_from_dir(test_dir)

assert isinstance(obj, AsyncStorageObject)
# Name should be directory name + .tar.gz
mock_async_client.objects.create.assert_awaited_once()
call_args = mock_async_client.objects.create.call_args
assert call_args[1]["name"] == "my_folder.tar.gz"
assert call_args[1]["content_type"] == "tgz"

@pytest.mark.asyncio
async def test_upload_from_dir_with_ttl(
self, mock_async_client: AsyncMock, object_view: MockObjectView, tmp_path: Path
) -> None:
"""Test upload_from_dir with TTL."""
from datetime import timedelta

mock_async_client.objects.create = AsyncMock(return_value=object_view)
mock_async_client.objects.complete = AsyncMock(return_value=object_view)

test_dir = tmp_path / "temp_dir"
test_dir.mkdir()
(test_dir / "file.txt").write_text("temporary content")

http_client = AsyncMock()
mock_response = create_mock_httpx_response()
http_client.put = AsyncMock(return_value=mock_response)
mock_async_client._client = http_client

client = AsyncStorageObjectOps(mock_async_client)
obj = await client.upload_from_dir(test_dir, ttl=timedelta(hours=2))

assert isinstance(obj, AsyncStorageObject)
mock_async_client.objects.create.assert_awaited_once()
call_args = mock_async_client.objects.create.call_args
# 2 hours = 7200 seconds = 7200000 milliseconds
assert call_args[1]["ttl_ms"] == 7200000

@pytest.mark.asyncio
async def test_upload_from_dir_empty_directory(
self, mock_async_client: AsyncMock, object_view: MockObjectView, tmp_path: Path
) -> None:
"""Test upload_from_dir with empty directory."""
mock_async_client.objects.create = AsyncMock(return_value=object_view)
mock_async_client.objects.complete = AsyncMock(return_value=object_view)

test_dir = tmp_path / "empty_dir"
test_dir.mkdir()

http_client = AsyncMock()
mock_response = create_mock_httpx_response()
http_client.put = AsyncMock(return_value=mock_response)
mock_async_client._client = http_client

client = AsyncStorageObjectOps(mock_async_client)
obj = await client.upload_from_dir(test_dir)

assert isinstance(obj, AsyncStorageObject)
assert obj.id == "obj_123"
mock_async_client.objects.create.assert_awaited_once()
http_client.put.assert_awaited_once()
mock_async_client.objects.complete.assert_awaited_once()

@pytest.mark.asyncio
async def test_upload_from_dir_with_string_path(
self, mock_async_client: AsyncMock, object_view: MockObjectView, tmp_path: Path
) -> None:
"""Test upload_from_dir with string path instead of Path object."""
mock_async_client.objects.create = AsyncMock(return_value=object_view)
mock_async_client.objects.complete = AsyncMock(return_value=object_view)

test_dir = tmp_path / "string_path_dir"
test_dir.mkdir()
(test_dir / "file.txt").write_text("content")

http_client = AsyncMock()
mock_response = create_mock_httpx_response()
http_client.put = AsyncMock(return_value=mock_response)
mock_async_client._client = http_client

client = AsyncStorageObjectOps(mock_async_client)
# Pass string path instead of Path object
obj = await client.upload_from_dir(str(test_dir))

assert isinstance(obj, AsyncStorageObject)
assert obj.id == "obj_123"
mock_async_client.objects.create.assert_awaited_once()


class TestAsyncRunloopSDK:
"""Tests for AsyncRunloopSDK class."""
Expand Down
Loading