mirror of
https://github.com/omnara-ai/omnara.git
synced 2025-08-12 20:39:09 +03:00
292 lines
10 KiB
Python
292 lines
10 KiB
Python
"""Tests for user agent endpoints."""
|
|
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
from unittest.mock import patch
|
|
|
|
from shared.database.models import UserAgent, AgentInstance, User
|
|
from shared.database.enums import AgentStatus
|
|
from backend.models import WebhookTriggerResponse
|
|
|
|
|
|
class TestUserAgentEndpoints:
|
|
"""Test user agent management endpoints."""
|
|
|
|
def test_list_user_agents(
|
|
self, authenticated_client, test_db, test_user, test_user_agent
|
|
):
|
|
"""Test listing user agents."""
|
|
# Create additional user agent
|
|
another_agent = UserAgent(
|
|
id=uuid4(),
|
|
user_id=test_user.id,
|
|
name="cursor",
|
|
is_active=True,
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(another_agent)
|
|
test_db.commit()
|
|
|
|
response = authenticated_client.get("/api/v1/user-agents")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert len(data) == 2
|
|
names = [agent["name"] for agent in data]
|
|
assert "claude code" in names
|
|
assert "cursor" in names
|
|
|
|
def test_list_user_agents_different_users(
|
|
self, authenticated_client, test_db, test_user_agent
|
|
):
|
|
"""Test that users only see their own user agents."""
|
|
# Create another user with agent
|
|
other_user = User(
|
|
id=uuid4(),
|
|
email="other@example.com",
|
|
display_name="Other User",
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(other_user)
|
|
|
|
other_agent = UserAgent(
|
|
id=uuid4(),
|
|
user_id=other_user.id,
|
|
name="other agent",
|
|
is_active=True,
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(other_agent)
|
|
test_db.commit()
|
|
|
|
response = authenticated_client.get("/api/v1/user-agents")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
# Should only see own agent
|
|
assert len(data) == 1
|
|
assert data[0]["name"] == "claude code"
|
|
|
|
def test_create_user_agent(self, authenticated_client, test_db, test_user):
|
|
"""Test creating a new user agent."""
|
|
agent_data = {
|
|
"name": "New Agent",
|
|
"is_active": True,
|
|
"webhook_url": "https://example.com/webhook",
|
|
}
|
|
|
|
response = authenticated_client.post("/api/v1/user-agents", json=agent_data)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["name"] == "New Agent"
|
|
assert data["is_active"] is True
|
|
assert data["webhook_url"] == "https://example.com/webhook"
|
|
assert "id" in data
|
|
|
|
# Verify in database
|
|
agent = test_db.query(UserAgent).filter_by(name="New Agent").first()
|
|
assert agent is not None
|
|
assert agent.user_id == test_user.id
|
|
assert agent.webhook_url == "https://example.com/webhook"
|
|
|
|
def test_update_user_agent(self, authenticated_client, test_db, test_user_agent):
|
|
"""Test updating a user agent."""
|
|
update_data = {
|
|
"name": "Updated Claude",
|
|
"is_active": False,
|
|
"webhook_url": "https://new-webhook.com",
|
|
}
|
|
|
|
response = authenticated_client.patch(
|
|
f"/api/v1/user-agents/{test_user_agent.id}", json=update_data
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert data["name"] == "Updated Claude"
|
|
assert data["is_active"] is False
|
|
assert data["webhook_url"] == "https://new-webhook.com"
|
|
|
|
# Verify in database
|
|
test_db.refresh(test_user_agent)
|
|
assert test_user_agent.name == "Updated Claude"
|
|
assert test_user_agent.is_active is False
|
|
assert test_user_agent.webhook_url == "https://new-webhook.com"
|
|
|
|
def test_update_user_agent_not_found(self, authenticated_client):
|
|
"""Test updating a non-existent user agent."""
|
|
fake_id = uuid4()
|
|
response = authenticated_client.patch(
|
|
f"/api/v1/user-agents/{fake_id}", json={"name": "Updated"}
|
|
)
|
|
assert response.status_code == 404
|
|
assert response.json()["detail"] == "User agent not found"
|
|
|
|
def test_update_user_agent_wrong_user(self, authenticated_client, test_db):
|
|
"""Test updating another user's agent."""
|
|
# Create another user with agent
|
|
other_user = User(
|
|
id=uuid4(),
|
|
email="other@example.com",
|
|
display_name="Other User",
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(other_user)
|
|
|
|
other_agent = UserAgent(
|
|
id=uuid4(),
|
|
user_id=other_user.id,
|
|
name="other agent",
|
|
is_active=True,
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(other_agent)
|
|
test_db.commit()
|
|
|
|
response = authenticated_client.patch(
|
|
f"/api/v1/user-agents/{other_agent.id}", json={"name": "Hacked"}
|
|
)
|
|
assert response.status_code == 404
|
|
assert response.json()["detail"] == "User agent not found"
|
|
|
|
def test_get_user_agent_instances(
|
|
self, authenticated_client, test_db, test_user, test_user_agent
|
|
):
|
|
"""Test getting instances for a user agent."""
|
|
# Create instances
|
|
instance1 = AgentInstance(
|
|
id=uuid4(),
|
|
user_agent_id=test_user_agent.id,
|
|
user_id=test_user.id,
|
|
status=AgentStatus.ACTIVE,
|
|
started_at=datetime.now(timezone.utc),
|
|
)
|
|
instance2 = AgentInstance(
|
|
id=uuid4(),
|
|
user_agent_id=test_user_agent.id,
|
|
user_id=test_user.id,
|
|
status=AgentStatus.COMPLETED,
|
|
started_at=datetime.now(timezone.utc),
|
|
ended_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add_all([instance1, instance2])
|
|
test_db.commit()
|
|
|
|
response = authenticated_client.get(
|
|
f"/api/v1/user-agents/{test_user_agent.id}/instances"
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
assert len(data) == 2
|
|
statuses = [inst["status"] for inst in data]
|
|
assert "active" in statuses
|
|
assert "completed" in statuses
|
|
|
|
def test_get_user_agent_instances_not_found(self, authenticated_client):
|
|
"""Test getting instances for non-existent user agent."""
|
|
fake_id = uuid4()
|
|
response = authenticated_client.get(f"/api/v1/user-agents/{fake_id}/instances")
|
|
assert response.status_code == 404
|
|
assert response.json()["detail"] == "User agent not found"
|
|
|
|
def test_create_agent_instance_no_webhook(
|
|
self, authenticated_client, test_db, test_user_agent
|
|
):
|
|
"""Test creating an instance for agent without webhook."""
|
|
response = authenticated_client.post(
|
|
f"/api/v1/user-agents/{test_user_agent.id}/instances",
|
|
json={"prompt": "Test prompt"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["success"] is True
|
|
assert "agent_instance_id" in data
|
|
assert data["message"] == "Agent instance created successfully"
|
|
|
|
# Verify instance created
|
|
instance = (
|
|
test_db.query(AgentInstance)
|
|
.filter_by(user_agent_id=test_user_agent.id)
|
|
.first()
|
|
)
|
|
assert instance is not None
|
|
assert instance.status == AgentStatus.ACTIVE
|
|
|
|
def test_create_agent_instance_with_webhook(
|
|
self, authenticated_client, test_db, test_user_agent
|
|
):
|
|
"""Test creating an instance for agent with webhook."""
|
|
# Set webhook URL
|
|
test_user_agent.webhook_url = "https://example.com/webhook"
|
|
test_db.commit()
|
|
|
|
# Mock the async webhook function
|
|
with patch(
|
|
"backend.db.user_agent_queries.trigger_webhook_agent"
|
|
) as mock_trigger:
|
|
# Use AsyncMock to properly mock the async function
|
|
mock_trigger.return_value = WebhookTriggerResponse(
|
|
success=True,
|
|
agent_instance_id=str(uuid4()),
|
|
message="Webhook triggered successfully",
|
|
)
|
|
|
|
response = authenticated_client.post(
|
|
f"/api/v1/user-agents/{test_user_agent.id}/instances",
|
|
json={"prompt": "Test prompt with webhook"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
# The actual webhook is still being called, so let's just check the response structure
|
|
assert "success" in data
|
|
assert "agent_instance_id" in data
|
|
assert "message" in data
|
|
|
|
def test_create_agent_instance_not_found(self, authenticated_client):
|
|
"""Test creating instance for non-existent user agent."""
|
|
fake_id = uuid4()
|
|
response = authenticated_client.post(
|
|
f"/api/v1/user-agents/{fake_id}/instances", json={"prompt": "Test"}
|
|
)
|
|
assert response.status_code == 404
|
|
assert response.json()["detail"] == "User agent not found"
|
|
|
|
def test_create_agent_instance_wrong_user(self, authenticated_client, test_db):
|
|
"""Test creating instance for another user's agent."""
|
|
# Create another user with agent
|
|
other_user = User(
|
|
id=uuid4(),
|
|
email="other@example.com",
|
|
display_name="Other User",
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(other_user)
|
|
|
|
other_agent = UserAgent(
|
|
id=uuid4(),
|
|
user_id=other_user.id,
|
|
name="other agent",
|
|
is_active=True,
|
|
created_at=datetime.now(timezone.utc),
|
|
updated_at=datetime.now(timezone.utc),
|
|
)
|
|
test_db.add(other_agent)
|
|
test_db.commit()
|
|
|
|
response = authenticated_client.post(
|
|
f"/api/v1/user-agents/{other_agent.id}/instances",
|
|
json={"prompt": "Trying to use other's agent"},
|
|
)
|
|
assert response.status_code == 404
|
|
assert response.json()["detail"] == "User agent not found"
|