KVM-privacy/services/kvm_agent/tests/test_mouse_ops.py

"""Unit tests for mouse_ops.py using mocked KVM and perception."""

import pytest
from unittest.mock import AsyncMock, MagicMock

from kvm_agent.mouse_ops import (
    click_element,
    click_dialog_button,
    close_window,
    show_desktop,
    safe_cleanup,
    _is_desktop,
    launch_from_taskbar,
    double_click,
    right_click,
    drag,
    scroll_in_area,
)
from kvm_agent.perception import SceneGraph, UIElement


# ── Helpers ────────────────────────────────────────────────────────


def _make_scene(elements=None, raw_text=""):
    return SceneGraph(elements=elements or [], raw_ocr_text=raw_text)


def _make_element(text, cx, cy, conf=0.9):
    w, h = 0.05, 0.02
    return UIElement(
        text=text,
        bbox=(cx - w, cy - h, cx + w, cy + h),
        confidence=conf,
    )


def _mock_kvm():
    kvm = AsyncMock()
    kvm.screenshot.return_value = b"fake_jpg"
    return kvm


def _mock_perception(elements=None, raw_text=""):
    perception = AsyncMock()
    scene = _make_scene(elements or [], raw_text)
    perception.perceive.return_value = scene
    return perception


# ═══════════════════════════════════════════════════════════════════
#  click_element
# ═══════════════════════════════════════════════════════════════════


class TestClickElement:

    @pytest.mark.asyncio
    async def test_found_and_clicked(self):
        kvm = _mock_kvm()
        perception = _mock_perception([_make_element("保存", 0.5, 0.3)])
        result = await click_element(kvm, "保存", perception)
        assert result is True
        kvm.mouse_move.assert_called_once()
        kvm.mouse_click.assert_called_once_with(0)

    @pytest.mark.asyncio
    async def test_not_found(self):
        kvm = _mock_kvm()
        perception = _mock_perception([])
        result = await click_element(kvm, "不存在", perception)
        assert result is False
        kvm.mouse_click.assert_not_called()

    @pytest.mark.asyncio
    async def test_skips_zero_bbox(self):
        """Elements with (0,0,0,0) bbox are skipped."""
        kvm = _mock_kvm()
        zero_el = UIElement(text="ghost", bbox=(0.0, 0.0, 0.0, 0.0), confidence=0.9)
        perception = _mock_perception([zero_el])
        result = await click_element(kvm, "ghost", perception)
        assert result is False
        kvm.mouse_click.assert_not_called()


# ═══════════════════════════════════════════════════════════════════
#  click_dialog_button
# ═══════════════════════════════════════════════════════════════════


class TestClickDialogButton:

    @pytest.mark.asyncio
    async def test_clicks_first_match(self):
        kvm = _mock_kvm()
        perception = _mock_perception([
            _make_element("确定", 0.4, 0.5),
            _make_element("取消", 0.6, 0.5),
        ])
        result = await click_dialog_button(kvm, perception, ["确定", "取消"])
        assert result is True
        # Should click 确定 (first in targets)
        args = kvm.mouse_move.call_args[0]
        assert abs(args[0] - 0.4) < 0.1

    @pytest.mark.asyncio
    async def test_no_match_returns_false(self):
        kvm = _mock_kvm()
        perception = _mock_perception([])
        result = await click_dialog_button(kvm, perception, ["OK"])
        assert result is False

    @pytest.mark.asyncio
    async def test_ignores_edge_elements(self):
        """Elements near edge (taskbar area) are skipped."""
        kvm = _mock_kvm()
        # Element at y=0.95 (taskbar area)
        edge_el = _make_element("取消", 0.5, 0.95)
        perception = _mock_perception([edge_el])
        result = await click_dialog_button(kvm, perception, ["取消"])
        assert result is False


# ═══════════════════════════════════════════════════════════════════
#  close_window
# ═══════════════════════════════════════════════════════════════════


class TestCloseWindow:

    @pytest.mark.asyncio
    async def test_strategy1_title_bar(self):
        """Strategy 1: find title bar text → X button at 0.98, same y."""
        kvm = _mock_kvm()
        # Title bar element in top region
        perception = _mock_perception([_make_element("记事本 - 无标题", 0.4, 0.02)])
        result = await close_window(kvm, perception)
        assert result is True
        # X button should be at x=0.98, close to title's y
        call_args = kvm.mouse_move.call_args[0]
        assert call_args[0] == pytest.approx(0.98)
        assert call_args[1] < 0.10

    @pytest.mark.asyncio
    async def test_strategy2_fallback(self):
        """Strategy 2: no title found → fallback to (0.98, 0.01)."""
        kvm = _mock_kvm()
        perception = _mock_perception([])  # No elements
        result = await close_window(kvm, perception)
        assert result is True
        call_args = kvm.mouse_move.call_args[0]
        assert call_args[0] == pytest.approx(0.98)
        assert call_args[1] == pytest.approx(0.01)

    @pytest.mark.asyncio
    async def test_ignores_non_title_elements(self):
        """Elements below 10% screen height are not considered title bar."""
        kvm = _mock_kvm()
        # Element at y=0.5 (middle of screen)
        perception = _mock_perception([_make_element("内容", 0.5, 0.5)])
        result = await close_window(kvm, perception)
        assert result is True
        # Should use fallback since element not in title bar region
        call_args = kvm.mouse_move.call_args[0]
        assert call_args[1] == pytest.approx(0.01)


# ═══════════════════════════════════════════════════════════════════
#  show_desktop
# ═══════════════════════════════════════════════════════════════════


class TestShowDesktop:

    @pytest.mark.asyncio
    async def test_clicks_bottom_right_corner(self):
        kvm = _mock_kvm()
        await show_desktop(kvm)
        kvm.mouse_move.assert_called_once_with(0.999, 0.999)
        kvm.mouse_click.assert_called_once_with(0)


# ═══════════════════════════════════════════════════════════════════
#  _is_desktop
# ═══════════════════════════════════════════════════════════════════


class TestIsDesktop:

    def test_desktop_true(self):
        scene = _make_scene(raw_text="搜索 开始 任务栏 回收站 2026")
        assert _is_desktop(scene) is True

    def test_desktop_false_with_app(self):
        scene = _make_scene(raw_text="搜索 开始 任务栏 notepad 文件")
        assert _is_desktop(scene) is False

    def test_no_desktop_indicators(self):
        scene = _make_scene(raw_text="random text without indicators")
        assert _is_desktop(scene) is False


# ═══════════════════════════════════════════════════════════════════
#  safe_cleanup
# ═══════════════════════════════════════════════════════════════════


class TestSafeCleanup:

    @pytest.mark.asyncio
    async def test_desktop_reached_immediately(self):
        kvm = _mock_kvm()
        perception = _mock_perception(raw_text="搜索 开始 任务栏 回收站")
        await safe_cleanup(kvm, perception, max_rounds=3)
        # Should not try to close anything
        assert kvm.mouse_move.call_count == 0

    @pytest.mark.asyncio
    async def test_shutdown_dialog_cancelled(self):
        kvm = _mock_kvm()
        shutdown_scene = _make_scene(
            [_make_element("取消", 0.5, 0.5)],
            raw_text="关机 希望计算机做什么",
        )
        desktop_scene = _make_scene(raw_text="搜索 开始 任务栏 回收站")
        # safe_cleanup calls screenshot+perceive per round,
        # click_dialog_button calls screenshot+perceive internally too
        perception = AsyncMock()
        perception.perceive.side_effect = [
            shutdown_scene,  # round 1: safe_cleanup detects shutdown
            shutdown_scene,  # round 1: click_dialog_button re-perceives
            desktop_scene,   # round 2: safe_cleanup sees desktop
        ]
        kvm.screenshot.side_effect = [b"f1", b"f2", b"f3"]

        await safe_cleanup(kvm, perception, max_rounds=3)
        # Should have clicked cancel button
        assert kvm.mouse_click.call_count >= 1

    @pytest.mark.asyncio
    async def test_max_rounds_triggers_show_desktop(self):
        kvm = _mock_kvm()
        # Always return an app window scene
        perception = _mock_perception(
            [_make_element("notepad", 0.5, 0.02)],
            raw_text="notepad 文件 编辑",
        )
        await safe_cleanup(kvm, perception, max_rounds=2)
        # After max rounds, should click show desktop at (0.999, 0.999)
        last_move = kvm.mouse_move.call_args_list[-1][0]
        assert last_move[0] == pytest.approx(0.999)
        assert last_move[1] == pytest.approx(0.999)


# ═══════════════════════════════════════════════════════════════════
#  Primitive operations
# ═══════════════════════════════════════════════════════════════════


class TestPrimitives:

    @pytest.mark.asyncio
    async def test_double_click(self):
        kvm = _mock_kvm()
        await double_click(kvm, 0.5, 0.5)
        assert kvm.mouse_click.call_count == 2

    @pytest.mark.asyncio
    async def test_right_click(self):
        kvm = _mock_kvm()
        await right_click(kvm, 0.3, 0.7)
        kvm.mouse_move.assert_called_once_with(0.3, 0.7)
        kvm.mouse_click.assert_called_once_with(2)

    @pytest.mark.asyncio
    async def test_drag(self):
        kvm = _mock_kvm()
        await drag(kvm, 0.1, 0.1, 0.9, 0.9, steps=3)
        # Should have multiple mouse_move calls (start + 3 steps + release)
        assert kvm.mouse_move.call_count >= 5

    @pytest.mark.asyncio
    async def test_scroll_in_area(self):
        kvm = _mock_kvm()
        await scroll_in_area(kvm, 0.5, 0.5, 3)
        kvm.mouse_move.assert_called_once_with(0.5, 0.5)
        kvm.mouse_scroll.assert_called_once_with(3)