Artemis-Eapil/tests/unit/sdk/test_agent_trace_finalization.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from artemis.sdk.agent import Agent
from artemis.context import DeviceContext, DevicePlatform
from artemis.runtime.device_lock import DeviceBusyError
from artemis.sdk.types.exceptions import AgentError


@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["failed", "completed", "cancelled"])
async def test_failure_snapshot_is_written_before_tracing_ends(tmp_path, status):
    agent = object.__new__(Agent)
    task = MagicMock(status=status)
    context = MagicMock(execution_target="computer", execution_setup=None)
    context.data_engine.base_dir = tmp_path
    with patch("artemis.runtime.diagnostics.write_runtime_diagnostics") as snapshot:
        await agent._finalize_tracing(task, context)
    if status == "failed":
        snapshot.assert_called_once_with(tmp_path, "failure", android=False)
    else:
        snapshot.assert_not_called()
    context.data_engine.end_session.assert_called_once_with(status=status)


@pytest.mark.asyncio
async def test_trace_finalization_failure_does_not_escape():
    agent = object.__new__(Agent)
    agent._finalize_tracing = AsyncMock(side_effect=ValueError("bad recording manifest"))
    task = MagicMock(status="completed")
    task.get_name.return_value = "successful-task"
    context = MagicMock()

    await agent._finalize_tracing_safely(task=task, context=context)

    agent._finalize_tracing.assert_awaited_once_with(task=task, context=context)


@pytest.mark.asyncio
async def test_agent_clean_disconnects_the_shared_ui_client():
    agent = object.__new__(Agent)
    agent._initialized = True
    agent._ui_adb_client = MagicMock()

    await agent.clean()

    agent._ui_adb_client.disconnect.assert_called_once_with()
    assert agent._initialized is False


@pytest.mark.asyncio
async def test_device_context_uses_adb_size_without_starting_ui_client():
    agent = object.__new__(Agent)
    agent._adb_client = MagicMock()
    agent._adb_client.device.return_value.window_size.return_value = (1080, 2424)
    agent._ui_adb_client = MagicMock()

    from artemis.context import DevicePlatform

    context = await agent._get_device_context("device-123", DevicePlatform.ANDROID)

    assert context.device_width == 1080
    assert context.device_height == 2424
    agent._ui_adb_client.get_screen_data.assert_not_called()


@pytest.mark.asyncio
async def test_task_that_never_acquires_queue_does_not_create_trace_session(tmp_path, monkeypatch):
    from artemis.runtime import task_result

    monkeypatch.setattr(task_result, "TRACES_PATH", tmp_path)
    agent = Agent()
    agent._initialized = True
    agent._device_context = DeviceContext(
        host_platform="WINDOWS",
        mobile_platform=DevicePlatform.ANDROID,
        device_id="device-123",
        device_width=1080,
        device_height=2424,
    )
    agent._adb_client = MagicMock()
    agent._ui_adb_client = MagicMock()
    agent._prepare_tracing = MagicMock()

    with patch(
        "artemis.sdk.agent.DeviceExecutionLock.acquire",
        side_effect=DeviceBusyError("queue cancelled"),
    ):
        with pytest.raises(DeviceBusyError, match="queue cancelled"):
            await agent.run_task(goal="queued task", profile="flash")

    agent._prepare_tracing.assert_not_called()
    outcomes = list(tmp_path.glob("*/run_outcome.json"))
    assert len(outcomes) == 1
    import json

    outcome = json.loads(outcomes[0].read_text(encoding="utf-8"))
    assert outcome["code"] == "initialization_failed"
    assert "queue cancelled" in outcome["error"]


@pytest.mark.asyncio
async def test_secure_keyguard_is_rejected_without_guessing_credentials():
    agent = object.__new__(Agent)
    agent._device_context = MagicMock(device_id="device-123")
    agent._adb_client = MagicMock()
    agent._adb_client.device.return_value.shell.return_value = (
        'User "Owner" (id=0): deviceLocked=1\nUser "Work" (id=10): deviceLocked=1\n'
    )

    with pytest.raises(AgentError, match="Unlock the device manually"):
        await agent._ensure_device_unlocked()


@pytest.mark.asyncio
async def test_locked_work_profile_does_not_block_unlocked_device_owner():
    agent = object.__new__(Agent)
    agent._device_context = MagicMock(device_id="device-123")
    agent._adb_client = MagicMock()
    agent._adb_client.device.return_value.shell.return_value = (
        'User "Owner" (id=0): deviceLocked=0\nUser "Work" (id=10): deviceLocked=1\n'
    )

    await agent._ensure_device_unlocked()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "profile,enabled,expected",
    [
        ("flash", False, False),
        ("flash", True, True),
        ("Flash", False, False),
        ("pro", False, True),
        ("pro", True, True),
    ],
)
async def test_agent_inherits_session_id_from_env_and_propagates_to_tracing(
    monkeypatch, profile, enabled, expected
):
    import uuid

    canonical_sid = "a1b2c3d4-e5f6-4a5b-8c9d-0e1f2a3b4c5d"
    monkeypatch.setenv("ARTEMIS_SESSION_ID", canonical_sid)
    monkeypatch.setenv("ARTEMIS_FLASH_FINAL_CHECK_ENABLED", str(enabled).lower())

    agent = Agent()
    assert agent._session_id == canonical_sid

    task = MagicMock()
    task.id = "fallback-id"
    task.get_name.return_value = "task-name"
    task.request.goal = "test goal"
    task.request.profile = profile
    task.request.task_name = "task-name"
    task.request.trace_path = "traces"

    context = MagicMock()
    context.device = None
    with patch("artemis.sdk.agent.DataEngine") as mock_data_engine_cls:
        mock_data_engine = MagicMock()
        mock_data_engine_cls.return_value = mock_data_engine

        agent._prepare_tracing(task=task, context=context)
        assert context.execution_setup.final_check_enabled is expected

        mock_data_engine.start_session.assert_called_once()
        call_kwargs = mock_data_engine.start_session.call_args.kwargs
        assert call_kwargs["session_id"] == uuid.UUID(canonical_sid)