kairo-code 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

kairo/backend/tests/e2e/agents/test_agent_commands.py

@@ -0,0 +1,456 @@
+"""
+End-to-end tests for Agent Commands functionality.
+
+Tests cover:
+- Issue restart command
+- Issue stop command
+- Commands dispatched via heartbeat
+- Command expiry
+"""
+
+import pytest
+import uuid
+from datetime import datetime, timedelta, UTC
+from httpx import AsyncClient
+from sqlalchemy import select
+
+from backend.models.agent import Agent, AgentCommand
+
+
+class TestIssueRestartCommand:
+    """Tests for POST /api/agents/{agent_id}/restart"""
+
+    @pytest.mark.asyncio
+    async def test_issue_restart_command_success(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should successfully issue a restart command."""
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["command_type"] == "restart"
+        assert data["status"] == "pending"
+        assert "id" in data
+        assert "created_at" in data
+        assert "expires_at" in data
+
+    @pytest.mark.asyncio
+    async def test_restart_increments_restart_count(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should increment restart_count when restart command is issued."""
+        original_count = test_agent.restart_count
+
+        await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.restart_count == original_count + 1
+        assert test_agent.last_restart_at is not None
+        assert test_agent.last_restart_reason == "user_requested"
+
+    @pytest.mark.asyncio
+    async def test_restart_creates_command_record(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should create a command record in the database."""
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        command_id = response.json()["id"]
+
+        stmt = select(AgentCommand).where(AgentCommand.id == command_id)
+        result = await db_session.execute(stmt)
+        command = result.scalar_one_or_none()
+
+        assert command is not None
+        assert command.command_type == "restart"
+        assert command.agent_id == test_agent.id
+        assert command.status == "pending"
+
+    @pytest.mark.asyncio
+    async def test_restart_agent_not_found(
+        self, client: AsyncClient, auth_headers: dict
+    ):
+        """Should return 404 for non-existent agent."""
+        fake_id = str(uuid.uuid4())
+
+        response = await client.post(
+            f"/api/agents/{fake_id}/restart", headers=auth_headers
+        )
+
+        assert response.status_code == 404
+
+
+class TestIssueStopCommand:
+    """Tests for POST /api/agents/{agent_id}/stop"""
+
+    @pytest.mark.asyncio
+    async def test_issue_stop_command_success(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should successfully issue a stop command."""
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/stop", headers=auth_headers
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["command_type"] == "stop"
+        assert data["status"] == "pending"
+
+    @pytest.mark.asyncio
+    async def test_stop_does_not_increment_restart_count(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should not increment restart_count for stop command."""
+        original_count = test_agent.restart_count
+
+        await client.post(
+            f"/api/agents/{test_agent.id}/stop", headers=auth_headers
+        )
+
+        await db_session.refresh(test_agent)
+
+        assert test_agent.restart_count == original_count
+
+    @pytest.mark.asyncio
+    async def test_stop_agent_not_found(
+        self, client: AsyncClient, auth_headers: dict
+    ):
+        """Should return 404 for non-existent agent."""
+        fake_id = str(uuid.uuid4())
+
+        response = await client.post(
+            f"/api/agents/{fake_id}/stop", headers=auth_headers
+        )
+
+        assert response.status_code == 404
+
+
+class TestIssueGenericCommand:
+    """Tests for POST /api/agents/{agent_id}/command"""
+
+    @pytest.mark.asyncio
+    async def test_issue_restart_via_generic_endpoint(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should issue restart command via generic endpoint."""
+        payload = {"command_type": "restart"}
+
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/command",
+            json=payload,
+            headers=auth_headers,
+        )
+
+        assert response.status_code == 200
+        assert response.json()["command_type"] == "restart"
+
+    @pytest.mark.asyncio
+    async def test_issue_stop_via_generic_endpoint(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should issue stop command via generic endpoint."""
+        payload = {"command_type": "stop"}
+
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/command",
+            json=payload,
+            headers=auth_headers,
+        )
+
+        assert response.status_code == 200
+        assert response.json()["command_type"] == "stop"
+
+    @pytest.mark.asyncio
+    async def test_issue_config_refresh_command(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should issue config_refresh command."""
+        payload = {"command_type": "config_refresh"}
+
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/command",
+            json=payload,
+            headers=auth_headers,
+        )
+
+        assert response.status_code == 200
+        assert response.json()["command_type"] == "config_refresh"
+
+    @pytest.mark.asyncio
+    async def test_issue_command_with_payload(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should issue command with custom payload."""
+        payload = {
+            "command_type": "config_refresh",
+            "payload": {"config_key": "max_tokens", "config_value": 4096},
+        }
+
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/command",
+            json=payload,
+            headers=auth_headers,
+        )
+
+        assert response.status_code == 200
+        command_id = response.json()["id"]
+
+        stmt = select(AgentCommand).where(AgentCommand.id == command_id)
+        result = await db_session.execute(stmt)
+        command = result.scalar_one_or_none()
+
+        assert command.payload == {"config_key": "max_tokens", "config_value": 4096}
+
+    @pytest.mark.asyncio
+    async def test_issue_invalid_command_type(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should reject invalid command type."""
+        payload = {"command_type": "invalid_command"}
+
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/command",
+            json=payload,
+            headers=auth_headers,
+        )
+
+        assert response.status_code == 422
+
+
+class TestCommandDispatchViaHeartbeat:
+    """Tests for command dispatch via heartbeat."""
+
+    @pytest.mark.asyncio
+    async def test_command_delivered_via_heartbeat(
+        self, client: AsyncClient, auth_headers: dict, api_key_headers: dict,
+        test_agent: Agent
+    ):
+        """Should deliver pending command via heartbeat response."""
+        # Issue a restart command
+        await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        # Send heartbeat
+        heartbeat_payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert len(data["commands"]) == 1
+        assert data["commands"][0]["type"] == "restart"
+
+    @pytest.mark.asyncio
+    async def test_command_has_signature(
+        self, client: AsyncClient, auth_headers: dict, api_key_headers: dict,
+        test_agent: Agent
+    ):
+        """Should include signature in delivered command."""
+        # Issue command
+        await client.post(
+            f"/api/agents/{test_agent.id}/stop", headers=auth_headers
+        )
+
+        # Send heartbeat
+        heartbeat_payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+
+        data = response.json()
+        command = data["commands"][0]
+        assert "signature" in command
+        assert len(command["signature"]) == 64  # SHA256 hex
+
+    @pytest.mark.asyncio
+    async def test_command_status_changes_to_dispatched(
+        self, client: AsyncClient, auth_headers: dict, api_key_headers: dict,
+        test_agent: Agent, db_session
+    ):
+        """Should change command status to 'dispatched' after delivery."""
+        # Issue command
+        cmd_response = await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+        command_id = cmd_response.json()["id"]
+
+        # Send heartbeat
+        heartbeat_payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+
+        # Check command status
+        stmt = select(AgentCommand).where(AgentCommand.id == command_id)
+        result = await db_session.execute(stmt)
+        command = result.scalar_one_or_none()
+
+        assert command.status == "dispatched"
+        assert command.dispatched_at is not None
+
+    @pytest.mark.asyncio
+    async def test_command_not_delivered_twice(
+        self, client: AsyncClient, auth_headers: dict, api_key_headers: dict,
+        test_agent: Agent
+    ):
+        """Should not deliver same command multiple times."""
+        # Issue command
+        await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        heartbeat_payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        # First heartbeat - should get command
+        response1 = await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+        assert len(response1.json()["commands"]) == 1
+
+        # Second heartbeat - should not get command again
+        response2 = await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+        assert len(response2.json()["commands"]) == 0
+
+    @pytest.mark.asyncio
+    async def test_multiple_commands_delivered_in_order(
+        self, client: AsyncClient, auth_headers: dict, api_key_headers: dict,
+        test_agent: Agent
+    ):
+        """Should deliver multiple pending commands in creation order."""
+        # Issue multiple commands
+        await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+        await client.post(
+            f"/api/agents/{test_agent.id}/stop", headers=auth_headers
+        )
+
+        heartbeat_payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+
+        data = response.json()
+        assert len(data["commands"]) == 2
+        assert data["commands"][0]["type"] == "restart"
+        assert data["commands"][1]["type"] == "stop"
+
+
+class TestCommandExpiry:
+    """Tests for command expiration."""
+
+    @pytest.mark.asyncio
+    async def test_command_has_expiry_time(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent
+    ):
+        """Should set expiry time on command (default ~5 minutes)."""
+        response = await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        data = response.json()
+        expires_at = datetime.fromisoformat(data["expires_at"].replace("Z", "+00:00"))
+        created_at = datetime.fromisoformat(data["created_at"].replace("Z", "+00:00"))
+
+        # Should expire within 3-10 minutes
+        assert expires_at > created_at + timedelta(minutes=3)
+        assert expires_at < created_at + timedelta(minutes=10)
+
+    @pytest.mark.asyncio
+    async def test_expired_command_not_delivered(
+        self, client: AsyncClient, api_key_headers: dict, test_agent: Agent,
+        agent_factory
+    ):
+        """Should not deliver expired commands."""
+        # Create an expired command
+        await agent_factory.create_command(
+            test_agent.id,
+            command_type="restart",
+            status="pending",
+            expires_at=datetime.now(UTC) - timedelta(minutes=10),
+        )
+
+        heartbeat_payload = {
+            "agent_id": test_agent.id,
+            "status": "online",
+        }
+
+        response = await client.post(
+            "/api/agents/heartbeat",
+            json=heartbeat_payload,
+            headers=api_key_headers,
+        )
+
+        data = response.json()
+        assert len(data["commands"]) == 0
+
+
+class TestCommandEventLogging:
+    """Tests for command event logging."""
+
+    @pytest.mark.asyncio
+    async def test_command_issued_event_logged(
+        self, client: AsyncClient, auth_headers: dict, test_agent: Agent, db_session
+    ):
+        """Should log command_issued event when command is created."""
+        from backend.models.agent import AgentEvent
+
+        await client.post(
+            f"/api/agents/{test_agent.id}/restart", headers=auth_headers
+        )
+
+        # Check for event
+        stmt = (
+            select(AgentEvent)
+            .where(AgentEvent.agent_id == test_agent.id)
+            .where(AgentEvent.event_type == "command_issued")
+        )
+        result = await db_session.execute(stmt)
+        event = result.scalar_one_or_none()
+
+        assert event is not None
+        assert event.event_data["command_type"] == "restart"