planar 0.5.0__py3-none-any.whl

planar/files/storage/test_s3.py

@@ -0,0 +1,299 @@
+import asyncio
+import os
+import uuid
+from contextlib import asynccontextmanager
+
+import boto3
+import botocore
+import botocore.client
+import botocore.exceptions
+import pytest
+
+from planar.files.storage.s3 import S3Storage
+
+pytestmark = pytest.mark.skipif(
+    os.getenv("PLANAR_TEST_S3_STORAGE", "0") != "1",
+    reason="S3 tests must be enabled via PLANAR_TEST_S3_STORAGE env var",
+)
+
+# --- Configuration for LocalStack/S3 Compatible Service ---
+
+S3_PORT = 4566
+# LocalStack S3 endpoint
+S3_ENDPOINT_URL = f"http://127.0.0.1:{S3_PORT}"
+# Dummy credentials for LocalStack (usually not strictly required)
+AWS_ACCESS_KEY_ID = "test"
+AWS_SECRET_ACCESS_KEY = "test"
+AWS_REGION = "us-east-1"
+# Generate a unique bucket name for each test run session
+SESSION_BUCKET_NAME = f"planar-test-bucket-{uuid.uuid4()}"
+
+
+@pytest.fixture()
+def s3_boto_client():  # Synchronous client
+    """Provides a boto3 S3 client for direct interaction (e.g., bucket creation)."""
+    client = boto3.client(
+        "s3",
+        endpoint_url=S3_ENDPOINT_URL,
+        aws_access_key_id=AWS_ACCESS_KEY_ID,
+        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
+        region_name=AWS_REGION,
+        config=botocore.client.Config(signature_version="s3v4"),
+    )
+    return client
+
+
+@pytest.fixture(autouse=True)
+async def ensure_s3_bucket(s3_boto_client):
+    """
+    Ensures the S3 bucket exists before tests run.
+    This runs automatically due to autouse=True.
+    """
+    print(f"Attempting to create bucket: {SESSION_BUCKET_NAME} at {S3_ENDPOINT_URL}")
+
+    create_kwargs = {
+        "Bucket": SESSION_BUCKET_NAME,
+    }
+
+    try:
+        await asyncio.to_thread(s3_boto_client.create_bucket, **create_kwargs)
+        print(f"Bucket {SESSION_BUCKET_NAME} created or confirmed existing.")
+    except botocore.exceptions.ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code")
+        if error_code in ("BucketAlreadyOwnedByYou", "BucketAlreadyExists"):
+            print(f"Bucket {SESSION_BUCKET_NAME} already exists.")
+        else:
+            pytest.fail(
+                f"Failed to create S3 bucket {SESSION_BUCKET_NAME} "
+                f"at {S3_ENDPOINT_URL}: {e}. Is LocalStack running?"
+            )
+    except Exception as e:
+        pytest.fail(
+            f"An unexpected error occurred during bucket creation for {SESSION_BUCKET_NAME}: {e}"
+        )
+
+    yield  # Tests run here
+
+
+@pytest.fixture
+async def s3_storage() -> S3Storage:
+    """Provides an instance of S3Storage configured for the test bucket."""
+    storage_instance = S3Storage(
+        bucket_name=SESSION_BUCKET_NAME,
+        endpoint_url=S3_ENDPOINT_URL,
+        access_key_id=AWS_ACCESS_KEY_ID,
+        secret_access_key=AWS_SECRET_ACCESS_KEY,
+        region=AWS_REGION,
+        presigned_url_ttl=60,
+    )
+    return storage_instance
+
+
+@asynccontextmanager
+async def cleanup_s3_object(storage: S3Storage, ref: str):
+    """Context manager to ensure an S3 object is deleted after use."""
+    try:
+        yield
+    finally:
+        try:
+            print(f"Cleaning up S3 object: {ref}")
+            await storage.delete(ref)
+        except FileNotFoundError:
+            print(f"Cleanup: S3 object {ref} already deleted or not found.")
+        except Exception as e:
+            print(f"Warning: Failed to cleanup S3 object {ref}: {e}")
+
+
+# --- Test Cases ---
+
+
+async def test_put_get_bytes(s3_storage: S3Storage):
+    """Test storing and retrieving raw bytes."""
+    test_data = b"some binary data \x00\xff for s3"
+    mime_type = "application/octet-stream"
+    ref = None
+    try:
+        ref = await s3_storage.put_bytes(test_data, mime_type=mime_type)
+        assert isinstance(ref, str)
+        # S3 keys don't have to be UUIDs, but our implementation generates them
+        try:
+            uuid.UUID(ref)
+        except ValueError:
+            pytest.fail(f"Returned ref '{ref}' is not a valid UUID string")
+
+        async with cleanup_s3_object(s3_storage, ref):
+            retrieved_data, retrieved_mime = await s3_storage.get_bytes(ref)
+
+            assert retrieved_data == test_data
+            # S3 might add charset or other params, check starts with
+            assert retrieved_mime is not None
+            assert retrieved_mime.startswith(mime_type)
+
+            # Check external URL (should be a presigned URL)
+            url = await s3_storage.external_url(ref)
+            assert url is not None
+            base_expected_url = f"{S3_ENDPOINT_URL}/{SESSION_BUCKET_NAME}/{ref}"
+            assert url.startswith(base_expected_url)
+            assert "X-Amz-Signature" in url
+            assert "X-Amz-Expires" in url
+
+    except Exception as e:
+        if ref:
+            await cleanup_s3_object(s3_storage, ref).__aexit__(None, None, None)
+        raise e
+
+
+async def test_put_get_string(s3_storage: S3Storage):
+    """Test storing and retrieving a string."""
+    test_string = "Hello, S3! This is a test string with Unicode: éàçü."
+    mime_type = "text/plain"
+    encoding = "utf-16"
+    ref = None
+    try:
+        # Store with explicit encoding and mime type
+        ref = await s3_storage.put_string(
+            test_string, encoding=encoding, mime_type=mime_type
+        )
+        expected_mime_type = f"{mime_type}; charset={encoding}"
+
+        async with cleanup_s3_object(s3_storage, ref):
+            retrieved_string, retrieved_mime = await s3_storage.get_string(
+                ref, encoding=encoding
+            )
+
+            assert retrieved_string == test_string
+            assert retrieved_mime == expected_mime_type
+
+    except Exception as e:
+        if ref:
+            await cleanup_s3_object(s3_storage, ref).__aexit__(None, None, None)
+        raise e
+
+    # Test default encoding (utf-8)
+    ref_utf8 = None
+    try:
+        ref_utf8 = await s3_storage.put_string(test_string, mime_type="text/html")
+        expected_mime_utf8 = "text/html; charset=utf-8"
+
+        async with cleanup_s3_object(s3_storage, ref_utf8):
+            retrieved_string_utf8, retrieved_mime_utf8 = await s3_storage.get_string(
+                ref_utf8
+            )
+            assert retrieved_string_utf8 == test_string
+            assert retrieved_mime_utf8 == expected_mime_utf8
+    except Exception as e:
+        if ref_utf8:
+            await cleanup_s3_object(s3_storage, ref_utf8).__aexit__(None, None, None)
+        raise e
+
+
+async def test_put_get_stream(s3_storage: S3Storage):
+    """Test storing data from an async generator stream."""
+    test_chunks = [b"s3_chunk1 ", b"s3_chunk2 ", b"s3_chunk3"]
+    full_data = b"".join(test_chunks)
+    mime_type = "image/jpeg"  # Different mime type for variety
+    ref = None
+
+    async def _test_stream():
+        for chunk in test_chunks:
+            yield chunk
+            await asyncio.sleep(0.01)  # Simulate async work
+
+    try:
+        ref = await s3_storage.put(_test_stream(), mime_type=mime_type)
+
+        async with cleanup_s3_object(s3_storage, ref):
+            stream, retrieved_mime = await s3_storage.get(ref)
+            retrieved_data = b""
+            async for chunk in stream:
+                retrieved_data += chunk
+
+            assert retrieved_data == full_data
+            assert retrieved_mime is not None
+            assert retrieved_mime.startswith(mime_type)
+    except Exception as e:
+        if ref:
+            await cleanup_s3_object(s3_storage, ref).__aexit__(None, None, None)
+        raise e
+
+
+async def test_put_no_mime_type(s3_storage: S3Storage):
+    """Test storing data without providing a mime type."""
+    test_data = b"s3 data without mime"
+    ref = None
+    try:
+        ref = await s3_storage.put_bytes(test_data)
+        async with cleanup_s3_object(s3_storage, ref):
+            retrieved_data, retrieved_mime = await s3_storage.get_bytes(ref)
+
+            assert retrieved_data == test_data
+            # S3 might assign a default mime type (like binary/octet-stream) or none
+            # Depending on the S3 provider, this might be None or a default
+            print(f"Retrieved mime type (no mime put): {retrieved_mime}")
+            # assert retrieved_mime is None or retrieved_mime == 'binary/octet-stream'
+            # For now, let's just check the data
+    except Exception as e:
+        if ref:
+            await cleanup_s3_object(s3_storage, ref).__aexit__(None, None, None)
+        raise e
+
+
+async def test_delete(s3_storage: S3Storage):
+    """Test deleting stored data."""
+    ref = await s3_storage.put_bytes(b"to be deleted from s3", mime_type="text/plain")
+
+    # Verify object exists before delete (optional, get raises if not found)
+    try:
+        _, _ = await s3_storage.get(ref)
+    except FileNotFoundError:
+        pytest.fail(f"Object {ref} should exist before deletion but was not found.")
+
+    # Delete the object
+    await s3_storage.delete(ref)
+
+    # Verify object is gone after delete
+    with pytest.raises(FileNotFoundError):
+        await s3_storage.get(ref)
+
+    # Deleting again should be idempotent (no error)
+    try:
+        await s3_storage.delete(ref)
+    except Exception as e:
+        pytest.fail(f"Deleting already deleted ref raised an exception: {e}")
+
+
+async def test_get_non_existent(s3_storage: S3Storage):
+    """Test getting a reference that does not exist."""
+    non_existent_ref = str(uuid.uuid4())
+    with pytest.raises(FileNotFoundError):
+        await s3_storage.get(non_existent_ref)
+
+
+async def test_delete_non_existent(s3_storage: S3Storage):
+    """Test deleting a reference that does not exist (should not raise error)."""
+    non_existent_ref = str(uuid.uuid4())
+    try:
+        await s3_storage.delete(non_existent_ref)
+    except Exception as e:
+        pytest.fail(f"Deleting non-existent ref raised an exception: {e}")
+
+
+async def test_external_url(s3_storage: S3Storage):
+    """Test that external_url returns a valid-looking presigned S3 object URL."""
+    ref = None
+    try:
+        ref = await s3_storage.put_bytes(b"some data for url test")
+        async with cleanup_s3_object(s3_storage, ref):
+            url = await s3_storage.external_url(ref)
+            assert url is not None
+            base_expected_url = f"{S3_ENDPOINT_URL}/{SESSION_BUCKET_NAME}/{ref}"
+            assert url.startswith(base_expected_url)
+            assert "X-Amz-Algorithm" in url
+            assert "X-Amz-Credential" in url
+            assert "X-Amz-Date" in url
+            assert "X-Amz-Expires" in url
+            assert "X-Amz-Signature" in url
+    except Exception as e:
+        if ref:
+            await cleanup_s3_object(s3_storage, ref).__aexit__(None, None, None)
+        raise e

planar/files/test_files.py

@@ -0,0 +1,283 @@
+"""
+Test file handling in Planar workflows.
+"""
+
+import uuid
+from pathlib import Path
+from typing import AsyncGenerator, cast
+
+import pytest
+from pydantic import BaseModel, Field
+from sqlmodel.ext.asyncio.session import AsyncSession
+
+from planar.app import PlanarApp
+from planar.config import sqlite_config
+from planar.files import PlanarFile
+from planar.files.models import PlanarFileMetadata
+from planar.files.storage.base import Storage
+from planar.workflows.decorators import workflow
+from planar.workflows.execution import execute
+from planar.workflows.models import Workflow
+
+app = PlanarApp(
+    config=sqlite_config(":memory:"),
+    title="Planar app for testing file workflows",
+    description="Testing",
+)
+
+
+@pytest.fixture(name="app")
+def app_fixture():
+    yield app
+
+
+@pytest.fixture
+async def planar_file(
+    storage: Storage,
+    session: AsyncSession,  # Change type hint
+) -> PlanarFile:
+    """Create a PlanarFile instance for testing."""
+    # Store test content
+    test_data = b"Test file content for workflow"
+    mime_type = "text/plain"
+
+    # Store the file and get a reference
+    storage_ref = await storage.put_bytes(test_data, mime_type=mime_type)
+
+    # Create and store the file metadata
+    file_metadata = PlanarFileMetadata(
+        filename="test_file.txt",
+        content_type=mime_type,
+        size=len(test_data),
+        storage_ref=storage_ref,
+    )
+    session.add(file_metadata)
+    await session.commit()
+    await session.refresh(file_metadata)
+
+    # Return a PlanarFile reference (not the full metadata)
+    return PlanarFile(
+        id=file_metadata.id,
+        filename=file_metadata.filename,
+        content_type=file_metadata.content_type,
+        size=file_metadata.size,
+    )
+
+
+# Define models for workflow testing
+class FileProcessingInput(BaseModel):
+    """Input model for a workflow that processes a file."""
+
+    title: str = Field(description="Title of the processing job")
+    file: PlanarFile = Field(description="The file to process")
+    max_chars: int = Field(description="Maximum characters to extract", default=100)
+
+
+class FileProcessingResult(BaseModel):
+    """Result model for a file processing workflow."""
+
+    title: str = Field(description="Title of the processing job")
+    characters: int = Field(description="Number of characters in the file")
+    content_preview: str = Field(description="Preview of the file content")
+    file_id: uuid.UUID = Field(description="ID of the processed file")
+
+
+async def test_workflow_with_planar_file(
+    session: AsyncSession,
+    planar_file: PlanarFile,
+):
+    """Test that a workflow can accept and process a PlanarFile input."""
+
+    @workflow()
+    async def file_processing_workflow(input_data: PlanarFile):
+        file_content = await input_data.get_content()
+        char_count = len(file_content)
+        content_str = file_content.decode("utf-8")
+        preview = content_str[:100]
+
+        # Return structured result
+        return FileProcessingResult(
+            title="Test File Processing",
+            characters=char_count,
+            content_preview=preview,
+            file_id=input_data.id,
+        )
+
+    wf = await file_processing_workflow.start(planar_file)
+    result = await execute(wf)
+
+    # Verify the result
+    assert isinstance(result, FileProcessingResult)
+    assert result.title == "Test File Processing"
+    assert result.characters == len(b"Test file content for workflow")
+    assert result.content_preview == "Test file content for workflow"
+    assert result.file_id == planar_file.id
+
+    # Verify the workflow completed successfully
+    updated_wf = await session.get(Workflow, wf.id)
+    assert updated_wf is not None
+    assert updated_wf.status == "succeeded"
+    assert updated_wf.args == [planar_file.model_dump(mode="json")]
+
+    # Verify that the result stored in the workflow is correct
+    workflow_result = cast(dict, updated_wf.result)
+    assert workflow_result["title"] == "Test File Processing"
+    assert workflow_result["characters"] == len(b"Test file content for workflow")
+    assert workflow_result["content_preview"] == "Test file content for workflow"
+    assert workflow_result["file_id"] == str(planar_file.id)
+
+
+TEST_BYTES = b"Test data for upload"
+TEST_FILENAME = "upload_test.txt"
+TEST_CONTENT_TYPE = "text/plain"
+TEST_SIZE = len(TEST_BYTES)
+DEFAULT_CONTENT_TYPE = "application/octet-stream"
+
+
+async def assert_upload_success(
+    uploaded_file: PlanarFile,
+    expected_filename: str,
+    expected_content: bytes,
+    expected_content_type: str,
+    expected_size: int,
+    session: AsyncSession,
+):
+    """Helper function to assert successful file upload."""
+    assert isinstance(uploaded_file, PlanarFile)
+    assert uploaded_file.filename == expected_filename
+    assert uploaded_file.content_type == expected_content_type
+    assert uploaded_file.size == expected_size
+    assert isinstance(uploaded_file.id, uuid.UUID)
+
+    # Verify database record
+    metadata = await session.get(PlanarFileMetadata, uploaded_file.id)
+    assert metadata is not None
+    assert metadata.filename == expected_filename
+    assert metadata.content_type == expected_content_type
+    assert metadata.size == expected_size
+    assert metadata.storage_ref is not None
+
+    # Verify stored content
+    retrieved_content = await uploaded_file.get_content()
+    assert retrieved_content == expected_content
+
+
+async def test_planar_file_upload_bytes(storage: Storage, session: AsyncSession):
+    """Test PlanarFile.upload with bytes content."""
+    uploaded_file = await PlanarFile.upload(
+        content=TEST_BYTES,
+        filename=TEST_FILENAME,
+        content_type="text/plain",
+        size=100,
+    )
+    await assert_upload_success(
+        uploaded_file,
+        TEST_FILENAME,
+        TEST_BYTES,
+        "text/plain",
+        100,
+        session,
+    )
+
+
+async def test_planar_file_upload_bytes_defaults(
+    storage: Storage, session: AsyncSession
+):
+    """Test PlanarFile.upload with bytes content using default size/type."""
+    uploaded_file = await PlanarFile.upload(content=TEST_BYTES, filename=TEST_FILENAME)
+    await assert_upload_success(
+        uploaded_file,
+        TEST_FILENAME,
+        TEST_BYTES,
+        DEFAULT_CONTENT_TYPE,  # Default type expected
+        TEST_SIZE,  # Size should be calculated
+        session,
+    )
+
+
+async def test_planar_file_upload_path(
+    storage: Storage, session: AsyncSession, tmp_path: Path
+):
+    """Test PlanarFile.upload with Path content."""
+    test_file = tmp_path / TEST_FILENAME
+    test_file.write_bytes(TEST_BYTES)
+
+    uploaded_file = await PlanarFile.upload(
+        content=test_file,
+        filename=TEST_FILENAME,
+        content_type=TEST_CONTENT_TYPE,
+        size=TEST_SIZE,
+    )
+    await assert_upload_success(
+        uploaded_file,
+        TEST_FILENAME,
+        TEST_BYTES,
+        TEST_CONTENT_TYPE,
+        TEST_SIZE,
+        session,
+    )
+
+
+async def test_planar_file_upload_path_defaults(
+    storage: Storage, session: AsyncSession, tmp_path: Path
+):
+    """Test PlanarFile.upload with Path content using default/inferred size/type."""
+    test_file = tmp_path / "another_test.json"  # Use different extension for inference
+    test_data = b'{"key": "value"}'
+    test_file.write_bytes(test_data)
+
+    uploaded_file = await PlanarFile.upload(
+        content=test_file,
+        filename="data.json",  # Ensure filename matches for inference
+    )
+    await assert_upload_success(
+        uploaded_file,
+        "data.json",
+        test_data,
+        "application/json",  # Inferred type expected
+        len(test_data),  # Size should be calculated
+        session,
+    )
+
+
+async def simple_byte_stream(
+    data: bytes, chunk_size: int = 10
+) -> AsyncGenerator[bytes, None]:
+    """Helper async generator for stream tests."""
+    for i in range(0, len(data), chunk_size):
+        yield data[i : i + chunk_size]
+
+
+async def test_planar_file_upload_stream(storage: Storage, session: AsyncSession):
+    """Test PlanarFile.upload with AsyncGenerator content."""
+    uploaded_file = await PlanarFile.upload(
+        content=simple_byte_stream(TEST_BYTES),
+        filename=TEST_FILENAME,
+        content_type=TEST_CONTENT_TYPE,
+        size=TEST_SIZE,
+    )
+    await assert_upload_success(
+        uploaded_file,
+        TEST_FILENAME,
+        TEST_BYTES,
+        TEST_CONTENT_TYPE,
+        TEST_SIZE,
+        session,
+    )
+
+
+async def test_planar_file_upload_stream_defaults(
+    storage: Storage, session: AsyncSession
+):
+    """Test PlanarFile.upload with AsyncGenerator content using default size/type."""
+    uploaded_file = await PlanarFile.upload(
+        content=simple_byte_stream(TEST_BYTES), filename=TEST_FILENAME
+    )
+    await assert_upload_success(
+        uploaded_file,
+        TEST_FILENAME,
+        TEST_BYTES,
+        DEFAULT_CONTENT_TYPE,  # Default type expected
+        -1,  # Size should be unknown (-1)
+        session,
+    )

planar/human/__init__.py

@@ -0,0 +1,2 @@
+from .human import Human, Timeout, complete_human_task, get_human_tasks  # noqa: F401
+from .models import HumanTask, HumanTaskResult  # noqa: F401