planar 0.9.3__py3-none-any.whl → 0.11.0__py3-none-any.whl

planar/test_sqlalchemy.py

@@ -1,193 +0,0 @@
-from uuid import uuid4
-
-import pytest
-from sqlalchemy.exc import DBAPIError
-from sqlalchemy.ext.asyncio import AsyncEngine
-from sqlmodel import col, insert, select, text
-
-from planar.db import PlanarSession, new_session
-from planar.modeling.orm.planar_base_entity import PlanarBaseEntity
-
-
-class SomeModel(PlanarBaseEntity, table=True):
-    name: str
-    value: int = 0
-
-
-async def test_run_transaction_success(tmp_db_engine):
-    uuid = uuid4()
-    uuid2 = uuid4()
-
-    async def transaction_func(session: PlanarSession):
-        await session.exec(
-            insert(SomeModel).values(id=uuid, name="test_item", value=42)  # type: ignore
-        )
-        await session.exec(
-            insert(SomeModel).values(id=uuid2, name="test_item2", value=42)  # type: ignore
-        )
-
-    async with new_session(tmp_db_engine) as session:
-        session.max_conflict_retries = 3
-        await session.run_transaction(transaction_func, session)
-
-    async with new_session(tmp_db_engine) as session:
-        items = (
-            await session.exec(select(SomeModel).order_by(col(SomeModel.name)))
-        ).all()
-        assert items == [
-            SomeModel(id=uuid, name="test_item", value=42),
-            SomeModel(id=uuid2, name="test_item2", value=42),
-        ]
-
-
-async def test_run_transaction_failure(tmp_db_engine):
-    async def transaction_func(session: PlanarSession):
-        await session.exec(insert(SomeModel).values(name="test_item", value=42))  # type: ignore
-        raise ValueError("Test error")
-        await session.exec(insert(SomeModel).values(name="test_item2", value=42))  # type: ignore
-
-    async with new_session(tmp_db_engine) as session:
-        with pytest.raises(ValueError, match="Test error"):
-            session.max_conflict_retries = 3
-            await session.run_transaction(transaction_func, session)
-
-    async with new_session(tmp_db_engine) as session:
-        items = (await session.exec(select(SomeModel))).all()
-        assert items == []
-
-
-async def test_run_transaction_concurrent_retry_success(tmp_db_engine):
-    attempts = 0
-    uuid = uuid4()
-    uuid2 = uuid4()
-
-    async def transaction_func(session: PlanarSession):
-        nonlocal attempts
-        await session.exec(
-            insert(SomeModel).values(id=uuid, name="test_item", value=42)  # type: ignore
-        )
-        if attempts == 0:
-            attempts += 1
-            raise DBAPIError(
-                "Test error", None, Exception("could not serialize access")
-            )
-        await session.exec(
-            insert(SomeModel).values(id=uuid2, name="test_item2", value=42)  # type: ignore
-        )
-
-    async with new_session(tmp_db_engine) as session:
-        session.max_conflict_retries = 1
-        await session.run_transaction(transaction_func, session)
-
-    async with new_session(tmp_db_engine) as session:
-        items = (
-            await session.exec(select(SomeModel).order_by(col(SomeModel.name)))
-        ).all()
-        assert items == [
-            SomeModel(id=uuid, name="test_item", value=42),
-            SomeModel(id=uuid2, name="test_item2", value=42),
-        ]
-
-
-async def test_run_transaction_concurrent_retry_failure(tmp_db_engine):
-    attempts = 0
-
-    async def transaction_func(session: PlanarSession):
-        nonlocal attempts
-        await session.exec(insert(SomeModel).values(name="test_item", value=42))  # type: ignore
-        if attempts < 2:
-            attempts += 1
-            raise DBAPIError(
-                "Test error", None, Exception("could not serialize access")
-            )
-        await session.exec(insert(SomeModel).values(name="test_item2", value=42))  # type: ignore
-
-    async with new_session(tmp_db_engine) as session:
-        with pytest.raises(DBAPIError, match="Test error"):
-            session.max_conflict_retries = 1
-            await session.run_transaction(transaction_func, session)
-
-    async with new_session(tmp_db_engine) as session:
-        items = (await session.exec(select(SomeModel))).all()
-        assert items == []
-
-
-async def test_serializable_transaction_failure_1(tmp_db_engine: AsyncEngine):
-    if tmp_db_engine.dialect.name != "postgresql":
-        return pytest.skip("Test requires PostgreSQL database")
-
-    async with new_session(tmp_db_engine) as setup_session:
-        # Setup: Insert initial data
-        async with setup_session.begin():
-            setup_session.add(SomeModel(id=uuid4(), name="initial", value=10))
-
-    # Create two sessions
-    async with (
-        new_session(tmp_db_engine) as session1,
-        new_session(tmp_db_engine) as session2,
-    ):
-        # Begin transactions in both sessions
-        await session1.begin()
-        await session2.begin()
-
-        # Set serializable isolation level
-        await session1.set_serializable_isolation()
-        await session2.set_serializable_isolation()
-
-        # Session 1: Read data
-        item1 = (
-            await session1.exec(select(SomeModel).where(SomeModel.name == "initial"))
-        ).one()
-        assert item1.value == 10
-
-        # Session 2: Read the same data
-        item2 = (
-            await session2.exec(select(SomeModel).where(SomeModel.name == "initial"))
-        ).one()
-        assert item2.value == 10
-
-        # Both sessions update the same row
-        item1.value += 5
-        item2.value += 3
-
-        # Session 1: Commit should succeed
-        await session1.commit()
-
-        # Session 2: Commit should fail with serialization error
-        with pytest.raises(DBAPIError, match="could not serialize access"):
-            await session2.commit()
-
-
-async def test_entity_schema_and_planar_schema_presence(tmp_db_engine: AsyncEngine):
-    table_name = SomeModel.__tablename__
-
-    async with new_session(tmp_db_engine) as session:
-        dialect = session.dialect.name
-
-        if dialect == "postgresql":
-            # Verify schemas include 'planar' and the default entity schema 'planar_entity'
-            res = await session.exec(
-                text("select schema_name from information_schema.schemata")  # type: ignore[arg-type]
-            )
-            schemas = {row[0] for row in res}
-            assert "planar" in schemas
-            assert "planar_entity" in schemas
-
-            # Verify SomeModel table is created in the entity schema
-            res = await session.exec(
-                text(
-                    "select table_schema from information_schema.tables where table_name = :tn"
-                ).bindparams(tn=table_name)  # type: ignore[arg-type]
-            )
-            table_schemas = {row[0] for row in res}
-            assert "planar_entity" in table_schemas
-            assert "public" not in table_schemas
-
-        else:
-            # SQLite: no schemas; ensure table exists
-            res = await session.exec(
-                text("select name from sqlite_master where type='table'")  # type: ignore[arg-type]
-            )
-            tables = {row[0] for row in res}
-            assert table_name in tables
-            assert not any(name.startswith("planar.") for name in tables)

planar/test_utils.py

@@ -1,105 +0,0 @@
-import asyncio
-import time
-from datetime import UTC, datetime
-
-import pytest
-
-from planar.utils import asyncify, utc_now
-
-
-async def test_asyncify_converts_sync_to_async():
-    """Test that asyncify correctly converts a synchronous function to an asynchronous one."""
-
-    def sync_function(x, y):
-        return x + y
-
-    async_function = asyncify(sync_function)
-
-    # Check that the function is now a coroutine function
-    assert asyncio.iscoroutinefunction(async_function)
-    assert not asyncio.iscoroutinefunction(sync_function)
-
-    # Check that it can be awaited
-    result = await async_function(5, 3)
-    assert result == 8
-
-
-async def test_asyncify_with_args_and_kwargs():
-    """Test that asyncify correctly passes positional and keyword arguments."""
-
-    def complex_function(a, b, c=0, d=0):
-        return a + b + c + d
-
-    async_function = asyncify(complex_function)
-
-    # Test with positional args only
-    result1 = await async_function(1, 2)
-    assert result1 == 3
-
-    # Test with positional and keyword args
-    result2 = await async_function(1, 2, c=3, d=4)
-    assert result2 == 10
-
-
-async def test_asyncify_preserves_exceptions():
-    """Test that asyncify preserves exceptions raised by the wrapped function."""
-
-    def failing_function():
-        raise ValueError("Expected error")
-
-    async_function = asyncify(failing_function)
-
-    with pytest.raises(ValueError, match="Expected error"):
-        await async_function()
-
-
-async def test_asyncify_non_blocking():
-    """Test that asyncify runs the function in a way that doesn't block the event loop."""
-    # This counter will be incremented by a task running concurrently with our slow function
-    counter = 0
-
-    @asyncify
-    def slow_function():
-        time.sleep(0.5)  # This would block the event loop if not run in executor
-        return counter
-
-    # This task will increment the counter while the slow function is running
-    async def increment_counter():
-        nonlocal counter
-        await asyncio.sleep(0.1)  # Short sleep to allow the slow function to start
-        for _ in range(10):
-            counter += 1
-            await asyncio.sleep(0.01)  # Short sleep to yield control
-
-    # Create increment task
-    task = asyncio.create_task(increment_counter())
-
-    # Run the async function
-    assert counter == 0
-    result = await slow_function()
-    # If the event loop was blocked, the counter would be 0
-    assert counter == 10
-    assert result == 10
-
-    await task
-
-
-def test_raises_when_applied_to_async_function():
-    """Test that asyncify raises an error when applied to an async function."""
-
-    async def async_function():
-        pass
-
-    with pytest.raises(ValueError, match="Function is already async"):
-        asyncify(async_function)
-
-
-def test_utc_now_returns_naive_utc():
-    """utc_now should return a naive datetime captured within two timestamps."""
-
-    before = datetime.now(UTC).replace(tzinfo=None)
-    result = utc_now()
-    after = datetime.now(UTC).replace(tzinfo=None)
-
-    assert result.tzinfo is None
-    assert before <= result <= after

planar/testing/test_memory_storage.py

@@ -1,143 +0,0 @@
-import asyncio
-import uuid
-
-import pytest
-
-from planar.testing.memory_storage import MemoryStorage
-
-
-@pytest.fixture
-async def storage() -> MemoryStorage:
-    """Provides an instance of MemoryStorage."""
-    return MemoryStorage()
-
-
-async def test_put_get_bytes(storage: MemoryStorage):
-    """Test storing and retrieving raw bytes."""
-    test_data = b"some binary data \x00\xff for memory"
-    mime_type = "application/octet-stream"
-
-    ref = await storage.put_bytes(test_data, mime_type=mime_type)
-    assert isinstance(ref, str)
-    try:
-        uuid.UUID(ref)  # Check if ref is a valid UUID string
-    except ValueError:
-        pytest.fail(f"Returned ref '{ref}' is not a valid UUID string")
-
-    retrieved_data, retrieved_mime = await storage.get_bytes(ref)
-
-    assert retrieved_data == test_data
-    assert retrieved_mime == mime_type
-
-    # Check internal state (optional)
-    assert ref in storage._blobs
-    assert ref in storage._mime_types
-    assert storage._blobs[ref] == test_data
-    assert storage._mime_types[ref] == mime_type
-
-
-async def test_put_get_string(storage: MemoryStorage):
-    """Test storing and retrieving a string."""
-    test_string = "Hello, memory! This is a test string with Unicode: éàçü."
-    mime_type = "text/plain"
-    encoding = "utf-16"
-
-    # Store with explicit encoding and mime type
-    ref = await storage.put_string(test_string, encoding=encoding, mime_type=mime_type)
-    expected_mime_type = f"{mime_type}; charset={encoding}"
-
-    retrieved_string, retrieved_mime = await storage.get_string(ref, encoding=encoding)
-
-    assert retrieved_string == test_string
-    assert retrieved_mime == expected_mime_type
-
-    # Test default encoding (utf-8)
-    ref_utf8 = await storage.put_string(test_string, mime_type="text/html")
-    expected_mime_utf8 = "text/html; charset=utf-8"
-    retrieved_string_utf8, retrieved_mime_utf8 = await storage.get_string(ref_utf8)
-    assert retrieved_string_utf8 == test_string
-    assert retrieved_mime_utf8 == expected_mime_utf8
-
-
-async def test_put_get_stream(storage: MemoryStorage):
-    """Test storing data from an async generator stream."""
-    test_chunks = [b"mem_chunk1 ", b"mem_chunk2 ", b"mem_chunk3"]
-    full_data = b"".join(test_chunks)
-    mime_type = "image/gif"
-
-    async def _test_stream():
-        for chunk in test_chunks:
-            yield chunk
-            await asyncio.sleep(0.01)  # Simulate async work
-
-    ref = await storage.put(_test_stream(), mime_type=mime_type)
-
-    stream, retrieved_mime = await storage.get(ref)
-    retrieved_data = b""
-    async for chunk in stream:
-        retrieved_data += chunk
-
-    assert retrieved_data == full_data
-    assert retrieved_mime == mime_type
-
-
-async def test_put_no_mime_type(storage: MemoryStorage):
-    """Test storing data without providing a mime type."""
-    test_data = b"memory data without mime"
-
-    ref = await storage.put_bytes(test_data)
-    retrieved_data, retrieved_mime = await storage.get_bytes(ref)
-
-    assert retrieved_data == test_data
-    assert retrieved_mime is None
-
-    # Check internal state
-    assert ref in storage._blobs
-    assert ref not in storage._mime_types
-
-
-async def test_delete(storage: MemoryStorage):
-    """Test deleting stored data."""
-    ref = await storage.put_bytes(b"to be deleted from memory", mime_type="text/plain")
-
-    # Verify data exists before delete (optional)
-    assert ref in storage._blobs
-    assert ref in storage._mime_types
-
-    await storage.delete(ref)
-
-    # Verify data is gone after delete
-    assert ref not in storage._blobs
-    assert ref not in storage._mime_types
-
-    # Try getting deleted ref
-    with pytest.raises(FileNotFoundError):
-        await storage.get(ref)
-
-
-async def test_get_non_existent(storage: MemoryStorage):
-    """Test getting a reference that does not exist."""
-    non_existent_ref = str(uuid.uuid4())
-    with pytest.raises(FileNotFoundError):
-        await storage.get(non_existent_ref)
-
-
-async def test_delete_non_existent(storage: MemoryStorage):
-    """Test deleting a reference that does not exist (should not raise error)."""
-    non_existent_ref = str(uuid.uuid4())
-    initial_blob_count = len(storage._blobs)
-    initial_mime_count = len(storage._mime_types)
-    try:
-        await storage.delete(non_existent_ref)
-        # Ensure no data was actually deleted
-        assert len(storage._blobs) == initial_blob_count
-        assert len(storage._mime_types) == initial_mime_count
-    except Exception as e:
-        pytest.fail(f"Deleting non-existent ref raised an exception: {e}")
-
-
-async def test_external_url(storage: MemoryStorage):
-    """Test that external_url returns None for memory storage."""
-    ref = await storage.put_bytes(b"some data for url test")
-    url = await storage.external_url(ref)
-    assert url is None

planar/workflows/test_concurrency_detection.py

@@ -1,120 +0,0 @@
-import asyncio
-import multiprocessing
-import multiprocessing.connection
-from multiprocessing.connection import Connection
-from uuid import UUID
-
-from planar.db import DatabaseManager, new_session
-from planar.session import engine_var, get_engine, session_var
-from planar.workflows.decorators import step, workflow
-from planar.workflows.exceptions import LockResourceFailed
-from planar.workflows.execution import (
-    _DEFAULT_LOCK_DURATION,
-    execute,
-)
-from planar.workflows.lock import lock_workflow
-from planar.workflows.models import Workflow, WorkflowStatus
-
-# bidirectional communication between the test process and the worker processes.
-conn: Connection
-
-
-@step(max_retries=0)
-async def dummy_step():
-    conn.send("waiting")
-    # Wait until "proceed" is received from the queue.
-    if conn.recv() != "proceed":
-        raise Exception('Expected "proceed"')
-    return "success"
-
-
-@workflow()
-async def dummy_workflow():
-    # Run the dummy step and return its result.
-    result = await dummy_step()
-    return result
-
-
-# copy of the resume_workflow function which allows more fine grained control from
-# the test process. This is fine because our goal is to test concurrency detection
-# implemented by the execute function.
-async def resume_with_semaphores(workflow_id: UUID):
-    engine = get_engine()
-    async with new_session(engine) as session:
-        tok = session_var.set(session)
-        try:
-            async with session.begin():
-                workflow = await session.get(Workflow, workflow_id)
-            if not workflow:
-                raise ValueError(f"Workflow {workflow_id} not found")
-            conn.send("ready")
-            # Wait until "start" is received on stdin.
-            if conn.recv() != "start":
-                raise Exception('Expected "start"')
-            async with lock_workflow(
-                workflow,
-                _DEFAULT_LOCK_DURATION,
-                retry_count=0,
-            ):
-                await execute(workflow)
-            conn.send("completed")
-        except LockResourceFailed:
-            conn.send("conflict")
-        finally:
-            session_var.reset(tok)
-
-
-# This worker function will be launched as a separate process.
-# It takes the workflow id, db_url and a multiprocess Pipe.
-def worker(wf_id: UUID, db_url: str, connection: Connection):
-    global conn
-    conn = connection
-    # Create a new engine for this process.
-    db_manager = DatabaseManager(db_url)
-    db_manager.connect()
-    engine = db_manager.get_engine()
-    engine_var.set(engine)
-    # Run the resume_with_semaphores coroutine.
-    # We use asyncio.run so that the worker’s event loop is independent.
-    asyncio.run(resume_with_semaphores(wf_id))
-
-
-async def test_concurrent_workflow_execution(tmp_db_url, tmp_db_engine):
-    async with new_session(tmp_db_engine) as session:
-        session_var.set(session)
-        wf: Workflow = await dummy_workflow.start()
-        wf_id = wf.id
-
-        # Launch two separate processes that attempt to resume the workflow concurrently.
-        p1_parent, p1_worker = multiprocessing.Pipe(duplex=True)
-        p2_parent, p2_worker = multiprocessing.Pipe(duplex=True)
-        p1 = multiprocessing.Process(target=worker, args=(wf_id, tmp_db_url, p1_worker))
-        p2 = multiprocessing.Process(target=worker, args=(wf_id, tmp_db_url, p2_worker))
-        p1.start()
-        p2.start()
-        # wait for both workers to fetch the workflow from the database.
-        assert p1_parent.recv() == "ready"
-        assert p2_parent.recv() == "ready"
-        # allow worker 1 to proceed.
-        p1_parent.send("start")
-        # wait for worker 1 to start the workflow and pause in the dummy step.
-        assert p1_parent.recv() == "waiting"
-        # allow worker 2 to proceed.
-        p2_parent.send("start")
-        # worker 2 should fail and will send a "conflict" message.
-        assert p2_parent.recv() == "conflict"
-        # allow worker 1 to proceed
-        p1_parent.send("proceed")
-        # worker 1 should complete the workflow and send a "completed" message.
-        assert p1_parent.recv() == "completed"
-        # cleanup workers
-        p1.join()
-        p2.join()
-
-        await session.refresh(wf)
-        assert wf, f"Workflow {wf_id} not found"
-        # Assert that the workflow completed successfully.
-        assert wf.status == WorkflowStatus.SUCCEEDED, (
-            f"Unexpected workflow status: {wf.status}"
-        )
-        assert wf.result == "success", f"Unexpected workflow result: {wf.result}"

planar/workflows/test_lock_timeout.py

@@ -1,140 +0,0 @@
-import asyncio
-from datetime import timedelta
-
-from planar.db import new_session
-from planar.session import session_var
-from planar.testing.synchronizable_tracer import SynchronizableTracer, TraceSpec
-from planar.utils import utc_now
-from planar.workflows.decorators import workflow
-from planar.workflows.execution import execute
-from planar.workflows.models import (
-    LockedResource,
-    Workflow,
-    WorkflowStatus,
-    workflow_exec_lock_key,
-)
-from planar.workflows.orchestrator import WorkflowOrchestrator
-from planar.workflows.step_core import Suspend, suspend
-from planar.workflows.tracing import tracer_var
-
-
-# Define a long-running workflow.
-@workflow()
-async def long_running_workflow():
-    # Simulate a long-running operation by sleeping 1 second.
-    await asyncio.sleep(1)
-    return "finished"
-
-
-async def test_lock_timer_extension(tmp_db_engine):
-    tracer = SynchronizableTracer()
-    tracer_var.set(tracer)
-    lock_acquired = tracer.instrument(
-        TraceSpec(function_name="lock_resource", message="commit")
-    )
-    lock_heartbeat = tracer.instrument(
-        TraceSpec(function_name="lock_heartbeat", message="commit")
-    )
-
-    async with new_session(tmp_db_engine) as session:
-        # This test verifies that when a workflow is executing, the heartbeat task
-        # (lock_heartbeat) extends the workflow's lock_until field. We run a
-        # long-running workflow (which sleeps for 1 second) with a short lock
-        # duration and heartbeat interval. While the workflow is running we query
-        # the stored workflow record and ensure that lock_until is updated
-        # (extended) by the heartbeat.
-
-        session_var.set(session)
-        # Start the workflow.
-        # Run workflow execution in the background with short durations so
-        # heartbeat kicks in quickly.
-        async with WorkflowOrchestrator.ensure_started(
-            lock_duration=timedelta(seconds=1)
-        ) as orchestrator:
-            wf: Workflow = await long_running_workflow.start()
-            wf_id = wf.id
-            lock_key = workflow_exec_lock_key(wf_id)
-
-            await lock_acquired.wait()
-
-            async with session.begin():
-                locked_resource = await session.get(LockedResource, lock_key)
-            assert locked_resource, "Expected a locked resource record"
-            lock_time_1 = locked_resource.lock_until
-            assert lock_time_1, "Lock time should be set"
-
-            # Wait a bit longer to allow another heartbeat cycle.
-            await lock_heartbeat.wait()
-            async with session.begin():
-                await session.refresh(locked_resource)
-            lock_time_2 = locked_resource.lock_until
-            assert lock_time_2, "Lock time should be set"
-
-            # The lock_time_2 should be later than lock_time_1 if the heartbeat is working.
-            assert lock_time_2 > lock_time_1, (
-                f"Expected lock_until to be extended by heartbeat: {lock_time_1} vs {lock_time_2}"
-            )
-
-            # Let the workflow finish.
-            await orchestrator.wait_for_completion(wf_id)
-
-            # Verify the workflow completed successfully.
-            await session.refresh(wf)
-            assert wf.status == WorkflowStatus.SUCCEEDED
-            assert wf.result == "finished"
-
-
-@workflow()
-async def crashed_worker_workflow():
-    # This workflow uses suspend() to simulate work that is paused.
-    # The first execution returns a Suspend object.
-    # When resumed it completes and returns "completed".
-    # First step: suspend (simulate waiting, e.g. because a worker had locked it).
-    await suspend(interval=timedelta(seconds=5))
-    # After the suspension it resumes here.
-    return "completed"
-
-
-async def test_orchestrator_resumes_crashed_worker(tmp_db_engine):
-    # This test simulates the scenario where a worker has “crashed” after
-    # locking a workflow. We start a workflow that suspends. Then we add a LockedResource
-    # record with an expired lock_until time to simulate a crashed
-
-    # Invoking the workflow_orchestrator (which polls for suspended workflows
-    # whose wakeup time is reached or that have expired locks) should cause the
-    # the workflow to be resumed. Finally, we verify that the workflow
-    # completes successfully. Start the workflow – its first execution will
-    # suspend.
-    async with new_session(tmp_db_engine) as session:
-        session_var.set(session)
-        wf = await crashed_worker_workflow.start()
-
-        result = await execute(wf)
-        assert isinstance(result, Suspend), (
-            "Expected the workflow to suspend on first execution."
-        )
-        # Simulate a crashed worker by directly changing the workflow record.
-        await session.refresh(wf)
-        # Force wakeup_at and lock_until to be in the past.
-        past_time = utc_now() - timedelta(seconds=1)
-        wf.wakeup_at = past_time
-        session.add(LockedResource(lock_key=f"workflow:{wf.id}", lock_until=past_time))
-        # Ensure it is marked as running, which would not normally be picked by
-        # the orchestrator
-        await session.commit()
-
-        # Now run the orchestrator, which polls for suspended workflows with
-        # wakeup_at <= now.
-        # We use a short poll interval.
-        async with WorkflowOrchestrator.ensure_started(
-            poll_interval=0.2
-        ) as orchestrator:
-            await orchestrator.wait_for_completion(wf.id)
-
-        await session.refresh(wf)
-        assert wf.status == WorkflowStatus.SUCCEEDED, (
-            f"Expected workflow status 'success' but got {wf.status}"
-        )
-        assert wf.result == "completed", (
-            f"Expected workflow result 'completed' but got {wf.result}"
-        )