78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple test script to verify agent run cancellation functionality.
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import time
|
|
import json
|
|
|
|
# Add the backend directory to the path
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'backend'))
|
|
|
|
# Mock the dependencies we don't need
|
|
class MockApp:
|
|
def __init__(self):
|
|
self.data_directory = "/tmp/horde_test"
|
|
self.tmp = MockTmpState()
|
|
self.event_listeners = []
|
|
|
|
def emit_event(self, name, data=None):
|
|
pass
|
|
|
|
class MockTmpState:
|
|
def __init__(self):
|
|
self.runs = {}
|
|
|
|
# Mock the generate_tmp_id function
|
|
def mock_generate_tmp_id(app):
|
|
return f"tmp-run-{int(time.time())}"
|
|
|
|
# Monkey patch the temporary module
|
|
import backend.src.temporary as temp_module
|
|
temp_module.generate_tmp_id = mock_generate_tmp_id
|
|
|
|
from backend.src.agent import AgentRun
|
|
|
|
async def test_cancellation():
|
|
print("Testing agent run cancellation...")
|
|
|
|
# Create a mock app
|
|
app = MockApp()
|
|
|
|
# Create an agent run
|
|
run = AgentRun.create(app, agent_id="test-agent", user_id="test-user")
|
|
print(f"Created agent run: {run.id}")
|
|
|
|
# Test that the run is not cancelled initially
|
|
assert not run.cancelled, "Run should not be cancelled initially"
|
|
print("✓ Run is not cancelled initially")
|
|
|
|
# Test cancellation
|
|
run.cancel()
|
|
assert run.cancelled, "Run should be cancelled after calling cancel()"
|
|
assert run.done, "Run should be marked as done after cancellation"
|
|
print("✓ Run is cancelled and marked as done")
|
|
|
|
# Test that calling cancel again doesn't cause issues
|
|
run.cancel()
|
|
print("✓ Multiple cancel calls work correctly")
|
|
|
|
# Test that the run has the expected attributes
|
|
assert hasattr(run, 'task'), "Run should have a task attribute"
|
|
assert hasattr(run, 'thread'), "Run should have a thread attribute"
|
|
assert hasattr(run, 'cancelled'), "Run should have a cancelled attribute"
|
|
print("✓ Run has all required cancellation attributes")
|
|
|
|
# Test serialization
|
|
serialized = run.serialize()
|
|
assert serialized['id'] == run.id, "Serialized run should have correct ID"
|
|
assert serialized['done'] == True, "Serialized run should be marked as done"
|
|
print("✓ Run serialization works correctly")
|
|
|
|
print("All tests passed! Cancellation functionality is working correctly.")
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
asyncio.run(test_cancellation()) |