"""Tests for agent log streaming functionality."""
from unittest.mock import AsyncMock, patch, MagicMock
import pytest
from preloop.agents.container import ContainerAgentExecutor
@pytest.fixture
def mock_docker():
"""Mock aiodocker Docker client."""
with patch("preloop.agents.container.aiodocker.Docker") as mock:
docker_instance = AsyncMock()
mock.return_value = docker_instance
# Mock containers API
mock_containers = AsyncMock()
docker_instance.containers = mock_containers
# Mock images API
mock_images = AsyncMock()
docker_instance.images = mock_images
yield docker_instance
@pytest.fixture
def container_executor():
"""Create a ContainerAgentExecutor instance."""
config = {
"max_iterations": 10,
"timeout": 3600,
}
return ContainerAgentExecutor(
agent_type="test-agent",
config=config,
image="test-image:latest",
use_kubernetes=False,
)
@pytest.fixture
def k8s_executor():
"""Create a Kubernetes ContainerAgentExecutor instance."""
config = {
"max_iterations": 10,
"timeout": 3600,
}
return ContainerAgentExecutor(
agent_type="test-agent",
config=config,
image="test-image:latest",
use_kubernetes=True,
)
class TestDockerLogStreaming:
"""Test Docker log streaming functionality."""
@pytest.mark.asyncio
async def test_stream_docker_logs_success(self, container_executor, mock_docker):
"""Test streaming logs from a Docker container."""
mock_container = AsyncMock()
# Mock streaming logs
async def mock_log_stream(*args, **kwargs):
# Simulate streaming log lines
for line in [b"Line 1\n", b"Line 2\n", b"Line 3\n"]:
yield line
mock_container.log = mock_log_stream
mock_docker.containers.get = AsyncMock(return_value=mock_container)
# Collect streamed lines
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
assert len(lines) == 3
assert lines[0] == "Line 1"
assert lines[1] == "Line 2"
assert lines[2] == "Line 3"
mock_docker.containers.get.assert_called_once_with("container-123")
@pytest.mark.asyncio
async def test_stream_docker_logs_empty_lines_filtered(
self, container_executor, mock_docker
):
"""Test that empty lines are filtered out."""
mock_container = AsyncMock()
async def mock_log_stream(*args, **kwargs):
for line in [b"Line 1\n", b"\n", b"Line 2\n", b"", b"Line 3\n"]:
yield line
mock_container.log = mock_log_stream
mock_docker.containers.get = AsyncMock(return_value=mock_container)
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
# Empty lines should be filtered
assert len(lines) == 3
assert lines == ["Line 1", "Line 2", "Line 3"]
@pytest.mark.asyncio
async def test_stream_docker_logs_string_format(
self, container_executor, mock_docker
):
"""Test that string logs (not bytes) are handled correctly."""
mock_container = AsyncMock()
async def mock_log_stream(*args, **kwargs):
# Some aiodocker versions return strings instead of bytes
for line in ["String line 1\n", "String line 2\n", "String line 3\n"]:
yield line
mock_container.log = mock_log_stream
mock_docker.containers.get = AsyncMock(return_value=mock_container)
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
assert len(lines) == 3
assert lines[0] == "String line 1"
assert lines[1] == "String line 2"
assert lines[2] == "String line 3"
@pytest.mark.asyncio
async def test_stream_docker_logs_unicode(self, container_executor, mock_docker):
"""Test streaming logs with unicode characters."""
mock_container = AsyncMock()
async def mock_log_stream(*args, **kwargs):
for line in [b"Hello \xf0\x9f\x91\x8b\n", b"Testing \xe2\x9c\x93\n"]:
yield line
mock_container.log = mock_log_stream
mock_docker.containers.get = AsyncMock(return_value=mock_container)
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
assert len(lines) == 2
assert "👋" in lines[0]
assert "✓" in lines[1]
@pytest.mark.asyncio
async def test_stream_docker_logs_handles_docker_error(
self, container_executor, mock_docker
):
"""Test that streaming handles Docker errors gracefully."""
from aiodocker.exceptions import DockerError
mock_docker.containers.get = AsyncMock(
side_effect=DockerError(500, {"message": "Container not found"})
)
# Should yield error message instead of crashing
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
assert len(lines) > 0
assert "[ERROR]" in lines[0]
assert "Failed to stream logs" in lines[0]
@pytest.mark.asyncio
async def test_stream_docker_logs_handles_unexpected_error(
self, container_executor, mock_docker
):
"""Test that streaming handles unexpected errors gracefully."""
mock_container = AsyncMock()
async def mock_log_stream(*args, **kwargs):
yield b"Line 1\n"
raise ValueError("Unexpected error")
mock_container.log = mock_log_stream
mock_docker.containers.get = AsyncMock(return_value=mock_container)
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
# Should get first line and then error
assert len(lines) >= 1
assert lines[0] == "Line 1"
# The error should be caught and logged
class TestKubernetesLogStreaming:
"""Test Kubernetes log streaming functionality."""
@pytest.mark.asyncio
async def test_stream_kubernetes_logs_success(self, k8s_executor):
"""Test streaming logs from a Kubernetes pod."""
with patch("preloop.agents.container.KUBERNETES_AVAILABLE", True):
with patch("preloop.agents.container.config") as mock_config:
with patch("preloop.agents.container.client") as mock_client:
# Mock K8s API initialization
mock_config.load_incluster_config = MagicMock()
mock_batch_api = AsyncMock()
mock_core_api = AsyncMock()
mock_client.BatchV1Api.return_value = mock_batch_api
mock_client.CoreV1Api.return_value = mock_core_api
# Mock pod listing
mock_pod = MagicMock()
mock_pod.metadata.name = "test-pod-123"
mock_pods_list = MagicMock()
mock_pods_list.items = [mock_pod]
mock_core_api.list_namespaced_pod = AsyncMock(
return_value=mock_pods_list
)
# Mock log streaming response
mock_response = AsyncMock()
async def mock_content():
for line in [b"K8s Line 1\n", b"K8s Line 2\n", b"K8s Line 3\n"]:
yield line
mock_response.content = mock_content()
mock_core_api.read_namespaced_pod_log = AsyncMock(
return_value=mock_response
)
# Collect streamed lines
lines = []
async for line in k8s_executor.stream_logs("job-name"):
lines.append(line)
assert len(lines) == 3
assert lines[0] == "K8s Line 1"
assert lines[1] == "K8s Line 2"
assert lines[2] == "K8s Line 3"
@pytest.mark.asyncio
async def test_stream_kubernetes_logs_no_pods(self, k8s_executor):
"""Test streaming when no pods are found."""
with patch("preloop.agents.container.KUBERNETES_AVAILABLE", True):
with patch("preloop.agents.container.config") as mock_config:
with patch("preloop.agents.container.client") as mock_client:
mock_config.load_incluster_config = MagicMock()
mock_batch_api = AsyncMock()
mock_core_api = AsyncMock()
mock_client.BatchV1Api.return_value = mock_batch_api
mock_client.CoreV1Api.return_value = mock_core_api
# Mock empty pod list
mock_pods_list = MagicMock()
mock_pods_list.items = []
mock_core_api.list_namespaced_pod = AsyncMock(
return_value=mock_pods_list
)
lines = []
async for line in k8s_executor.stream_logs("job-name"):
lines.append(line)
assert len(lines) == 1
assert "[WARN]" in lines[0]
assert "No pods found" in lines[0]
@pytest.mark.asyncio
async def test_stream_kubernetes_logs_handles_api_exception(self, k8s_executor):
"""Test that streaming handles Kubernetes API exceptions."""
from kubernetes_asyncio.client.rest import ApiException
with patch("preloop.agents.container.KUBERNETES_AVAILABLE", True):
with patch("preloop.agents.container.config") as mock_config:
with patch("preloop.agents.container.client") as mock_client:
mock_config.load_incluster_config = MagicMock()
mock_batch_api = AsyncMock()
mock_core_api = AsyncMock()
mock_client.BatchV1Api.return_value = mock_batch_api
mock_client.CoreV1Api.return_value = mock_core_api
# Mock API exception with proper ApiException class
api_error = ApiException(status=404, reason="Not Found")
mock_core_api.list_namespaced_pod = AsyncMock(side_effect=api_error)
lines = []
async for line in k8s_executor.stream_logs("job-name"):
lines.append(line)
assert len(lines) > 0
assert "[WARN]" in lines[0] or "[ERROR]" in lines[0]
class TestStreamLogsUnified:
"""Test the unified stream_logs interface."""
@pytest.mark.asyncio
async def test_stream_logs_routes_to_docker(self, container_executor, mock_docker):
"""Test that stream_logs routes to Docker implementation."""
mock_container = AsyncMock()
async def mock_log_stream(*args, **kwargs):
yield b"Docker log\n"
mock_container.log = mock_log_stream
mock_docker.containers.get = AsyncMock(return_value=mock_container)
lines = []
async for line in container_executor.stream_logs("container-123"):
lines.append(line)
assert len(lines) == 1
assert lines[0] == "Docker log"
@pytest.mark.asyncio
async def test_stream_logs_routes_to_kubernetes(self, k8s_executor):
"""Test that stream_logs routes to Kubernetes implementation."""
with patch("preloop.agents.container.KUBERNETES_AVAILABLE", True):
with patch("preloop.agents.container.config") as mock_config:
with patch("preloop.agents.container.client") as mock_client:
mock_config.load_incluster_config = MagicMock()
mock_batch_api = AsyncMock()
mock_core_api = AsyncMock()
mock_client.BatchV1Api.return_value = mock_batch_api
mock_client.CoreV1Api.return_value = mock_core_api
mock_pod = MagicMock()
mock_pod.metadata.name = "test-pod"
mock_pods_list = MagicMock()
mock_pods_list.items = [mock_pod]
mock_core_api.list_namespaced_pod = AsyncMock(
return_value=mock_pods_list
)
mock_response = AsyncMock()
async def mock_content():
yield b"K8s log\n"
mock_response.content = mock_content()
mock_core_api.read_namespaced_pod_log = AsyncMock(
return_value=mock_response
)
lines = []
async for line in k8s_executor.stream_logs("job-name"):
lines.append(line)
assert len(lines) == 1
assert lines[0] == "K8s log"
