Skip to content

Commit

Permalink
feat(agents-api): Add test for foreach step / wait for input
Browse files Browse the repository at this point in the history
Signed-off-by: Diwank Singh Tomer <[email protected]>
  • Loading branch information
creatorrr committed Oct 10, 2024
1 parent ac9b7c1 commit bf2552d
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 8 deletions.
82 changes: 76 additions & 6 deletions agents-api/tests/test_execution_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -671,8 +671,7 @@ async def _(
assert result == expected_output


# FIXME: This test is not working. It gets stuck
# @test("workflow: wait for input step start")
@test("workflow: wait for input step start")
async def _(
client=cozo_client,
developer_id=test_developer_id,
Expand Down Expand Up @@ -710,7 +709,12 @@ async def _(
mock_run_task_execution_workflow.assert_called_once()

# Let it run for a bit
await asyncio.sleep(3)
result_coroutine = handle.result()
task = asyncio.create_task(result_coroutine)
try:
await asyncio.wait_for(task, timeout=3)
except asyncio.TimeoutError:
task.cancel()

# Get the history
history = await handle.fetch_history()
Expand All @@ -728,12 +732,78 @@ async def _(
activity for activity in activities_scheduled if activity
]

future = handle.result()
await future

assert "wait_for_input_step" in activities_scheduled


@test("workflow: foreach wait for input step start")
async def _(
client=cozo_client,
developer_id=test_developer_id,
agent=test_agent,
):
data = CreateExecutionRequest(input={"test": "input"})

task = create_task(
developer_id=developer_id,
agent_id=agent.id,
data=CreateTaskRequest(
**{
"name": "test task",
"description": "test task about",
"input_schema": {"type": "object", "additionalProperties": True},
"main": [
{
"foreach": {
"in": "'a b c'.split()",
"do": {"wait_for_input": {"info": {"hi": '"bye"'}}},
},
},
],
}
),
client=client,
)

async with patch_testing_temporal() as (_, mock_run_task_execution_workflow):
execution, handle = await start_execution(
developer_id=developer_id,
task_id=task.id,
data=data,
client=client,
)

assert handle is not None
assert execution.task_id == task.id
assert execution.input == data.input
mock_run_task_execution_workflow.assert_called_once()

# Let it run for a bit
result_coroutine = handle.result()
task = asyncio.create_task(result_coroutine)
try:
await asyncio.wait_for(task, timeout=3)
except asyncio.TimeoutError:
task.cancel()

# Get the history
history = await handle.fetch_history()
events = [MessageToDict(e) for e in history.events]
assert len(events) > 0

activities_scheduled = [
event.get("activityTaskScheduledEventAttributes", {})
.get("activityType", {})
.get("name")
for event in events
if "ACTIVITY_TASK_SCHEDULED" in event["eventType"]
]
activities_scheduled = [
activity for activity in activities_scheduled if activity
]

assert "for_each_step" in activities_scheduled


@test("workflow: if-else step")
async def _(
client=cozo_client,
Expand Down
18 changes: 16 additions & 2 deletions scheduler/docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ services:
- POSTGRES_PASSWORD=${TEMPORAL_POSTGRES_PASSWORD}
healthcheck:
test: [ "CMD-SHELL", "pg_isready -d ${TEMPORAL_POSTGRES_DB:-temporal} -U ${TEMPORAL_POSTGRES_USER:-temporal}" ]

interval: 1s
timeout: 5s
retries: 10
Expand All @@ -63,8 +63,22 @@ services:
profiles:
- temporal-ui
environment:
# See: https://github.com/temporalio/ui-server/blob/main/docker/config-template.yaml
- TEMPORAL_ADDRESS=${TEMPORAL_ADDRESS:-temporal:7233}
- TEMPORAL_CORS_ORIGINS=${TEMPORAL_CORS_ORIGINS:-http://localhost:3000}
# Note: Not setting it enables all origins
# - TEMPORAL_CORS_ORIGINS=${TEMPORAL_CORS_ORIGINS:-http://localhost:3000}
- TEMPORAL_CODEC_ENDPOINT=http://localhost/api/temporal
- TEMPORAL_UI_ENABLED=true
- TEMPORAL_FEEDBACK_URL=https://github.com/julep-ai/julep
- TEMPORAL_NOTIFY_ON_NEW_VERSION=false
- TEMPORAL_CSRF_COOKIE_INSECURE=true
# - TEMPORAL_HIDE_LOGS=true
# - TEMPORAL_BANNER_TEXT=gagagaga
# - TEMPORAL_DISABLE_WRITE_ACTIONS=true
# - TEMPORAL_HIDE_WORKFLOW_QUERY_ERRORS=true
# - TEMPORAL_REFRESH_WORKFLOW_COUNTS_DISABLED=true
- TEMPORAL_OPEN_API_ENABLED=true

ports:
- 9000:8080 # Since 8080 is already used by agents-api

Expand Down

0 comments on commit bf2552d

Please sign in to comment.