392 lines
14 KiB
Python
392 lines
14 KiB
Python
"""Recorder lifecycle and archive limits use fake processes, never a live desktop."""
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from artemis.config.agent import RecordingConfig
|
|
from artemis.plugins.models import EvidenceSource
|
|
from artemis.plugins import recording
|
|
from artemis.plugins.recording import PluginRecordings, build_capture_command, capture_rectangle
|
|
|
|
|
|
def source(**updates):
|
|
return EvidenceSource(
|
|
**{
|
|
"kind": "computer",
|
|
"target_id": "123",
|
|
"label": "Fixture window",
|
|
"width": 800,
|
|
"height": 600,
|
|
"window": {"left": 40, "top": 50, "width": 800, "height": 600},
|
|
"screen": {"left": -1920, "top": 0, "width": 1920, "height": 1080},
|
|
**updates,
|
|
}
|
|
)
|
|
|
|
|
|
def context(root, *, enabled=True, **config):
|
|
return SimpleNamespace(
|
|
data_engine=SimpleNamespace(base_dir=root, session_start_time=100),
|
|
agent_config=SimpleNamespace(recording=RecordingConfig(**config)),
|
|
execution_setup=SimpleNamespace(
|
|
record_video_enabled=enabled, video_recording_tools_enabled=False
|
|
),
|
|
)
|
|
|
|
|
|
class FakeRecorderProcess:
|
|
def __init__(self, *, produces_frames=True):
|
|
self.pid = 123456
|
|
self.returncode = None if produces_frames else 1
|
|
self.stdout = asyncio.StreamReader()
|
|
if produces_frames:
|
|
self.stdout.feed_data(b"frame=1\n")
|
|
self.stdout.feed_eof()
|
|
self.stdin = SimpleNamespace(write=self._quit, drain=AsyncMock())
|
|
self.killed = False
|
|
|
|
def _quit(self, _data):
|
|
self.returncode = 0
|
|
|
|
async def wait(self):
|
|
return self.returncode
|
|
|
|
def kill(self):
|
|
self.killed = True
|
|
self.returncode = -9
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_capture(monkeypatch):
|
|
state = SimpleNamespace(commands=[], processes=[], fail_hevc=False)
|
|
|
|
async def spawn(*command, **_kwargs):
|
|
state.commands.append(command)
|
|
encoder = command[command.index("-c:v") + 1]
|
|
process = FakeRecorderProcess(
|
|
produces_frames=not (state.fail_hevc and encoder.startswith("hevc"))
|
|
)
|
|
state.processes.append(process)
|
|
Path(command[-1]).write_bytes(b"fake mp4")
|
|
return process
|
|
|
|
local_asyncio = SimpleNamespace(**vars(asyncio))
|
|
local_asyncio.create_subprocess_exec = spawn
|
|
monkeypatch.setattr(recording, "asyncio", local_asyncio)
|
|
monkeypatch.setattr(recording, "os", SimpleNamespace(name="nt", getpid=lambda: 123))
|
|
monkeypatch.setattr(
|
|
recording,
|
|
"CaptureProcessJob",
|
|
lambda: SimpleNamespace(assign=lambda _pid: None, close=lambda: None),
|
|
)
|
|
monkeypatch.setattr(
|
|
recording, "shutil", SimpleNamespace(disk_usage=lambda _path: SimpleNamespace(free=10**10))
|
|
)
|
|
monkeypatch.setattr(recording, "time", SimpleNamespace(time=lambda: 123.0))
|
|
monkeypatch.setattr(recording, "get_ffmpeg_path", lambda: "fake-ffmpeg")
|
|
monkeypatch.setattr(
|
|
recording,
|
|
"probe_video_segment",
|
|
AsyncMock(return_value={"duration": 2.5, "width": 1280, "height": 720}),
|
|
)
|
|
return state
|
|
|
|
|
|
def argument(command, option):
|
|
return command[command.index(option) + 1]
|
|
|
|
|
|
def test_archive_command_keeps_physical_capture_and_limits_encoded_output(tmp_path):
|
|
command = build_capture_command(
|
|
"ffmpeg", source(), RecordingConfig(), tmp_path / "video.mp4", "libx264"
|
|
)
|
|
assert argument(command, "-offset_x") == "-1920"
|
|
assert argument(command, "-video_size") == "1920x1080"
|
|
assert argument(command, "-vf") == "scale=1280:720"
|
|
assert argument(command, "-framerate") == "15"
|
|
assert argument(command, "-b:v") == "1500k"
|
|
assert argument(command, "-g") == "30"
|
|
assert "-an" in command
|
|
|
|
|
|
def test_small_window_is_not_upscaled_and_hwnd_capture_is_explicit(tmp_path):
|
|
command = build_capture_command(
|
|
"ffmpeg",
|
|
source(),
|
|
RecordingConfig(computer_scope="window", codec="h265"),
|
|
tmp_path / "video.mp4",
|
|
"hevc_nvenc",
|
|
)
|
|
assert argument(command, "-i") == "hwnd=123"
|
|
assert argument(command, "-vf") == "scale=800:600"
|
|
assert argument(command, "-b:v") == "1000k"
|
|
assert argument(command, "-tag:v") == "hvc1"
|
|
assert "-offset_x" not in command
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"area",
|
|
[
|
|
None,
|
|
{"left": 0},
|
|
{"left": 0, "top": 0, "width": 0, "height": 600},
|
|
{"left": 65537, "top": 0, "width": 800, "height": 600},
|
|
],
|
|
)
|
|
def test_invalid_capture_geometry_is_not_sent_to_ffmpeg(area):
|
|
with pytest.raises(ValueError):
|
|
capture_rectangle(source(screen=area), "screen")
|
|
|
|
|
|
def test_window_recording_rejects_non_hwnd_targets(tmp_path):
|
|
with pytest.raises(ValueError, match="HWND"):
|
|
build_capture_command(
|
|
"ffmpeg",
|
|
source(target_id="window-title"),
|
|
RecordingConfig(computer_scope="window"),
|
|
tmp_path / "video.mp4",
|
|
"libx264",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disabled_recording_does_not_spawn_or_create_evidence(tmp_path, fake_capture):
|
|
recorder = PluginRecordings(context(tmp_path, enabled=False))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
await recorder.stop()
|
|
assert fake_capture.commands == []
|
|
assert not (tmp_path / "evidence").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"codec,hevc_playback,expected_attempts",
|
|
[("auto", False, 1), ("auto", True, 4), ("h265", False, 4)],
|
|
)
|
|
async def test_hardware_hevc_failure_falls_back_to_compatible_archive(
|
|
tmp_path, fake_capture, codec, hevc_playback, expected_attempts
|
|
):
|
|
fake_capture.fail_hevc = True
|
|
recorder = PluginRecordings(context(tmp_path, codec=codec, hevc_playback=hevc_playback))
|
|
try:
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
assert len(fake_capture.commands) == expected_attempts
|
|
assert argument(fake_capture.commands[-1], "-c:v") == "libx264"
|
|
assert len(recorder.active) == 1
|
|
finally:
|
|
await recorder.stop()
|
|
track = next(iter(recorder.sources.values()))
|
|
assert track["status"] == "ready"
|
|
assert len(track["segments"]) == 1
|
|
assert track["segments"][0]["codec"] == "h264"
|
|
assert track["segments"][0]["duration_ms"] == 2500
|
|
assert all(process.returncode is not None for process in fake_capture.processes)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_same_screen_reuses_recording_and_keeps_both_target_ids(tmp_path, fake_capture):
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
try:
|
|
await recorder.observe("window-1", "visual-test", source())
|
|
await recorder.observe("window-2", "visual-test", source(target_id="456"))
|
|
assert len(fake_capture.commands) == 1
|
|
finally:
|
|
await recorder.stop()
|
|
assert next(iter(recorder.sources.values()))["target_ids"] == ["window-1", "window-2"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_low_disk_space_records_failure_without_starting_capture(
|
|
tmp_path, fake_capture, monkeypatch
|
|
):
|
|
monkeypatch.setattr(recording.shutil, "disk_usage", lambda _path: SimpleNamespace(free=1))
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
await recorder.stop()
|
|
assert fake_capture.commands == []
|
|
track = next(iter(recorder.sources.values()))
|
|
assert track["status"] == "failed"
|
|
assert "磁盘" in track["message"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_write_failure_does_not_leave_capture_processes(
|
|
tmp_path, fake_capture, monkeypatch
|
|
):
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
|
|
def unavailable():
|
|
raise OSError("disk full")
|
|
|
|
monkeypatch.setattr(recorder, "_save", unavailable)
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
await recorder.stop()
|
|
assert not recorder.active
|
|
assert all(process.returncode is not None for process in fake_capture.processes)
|
|
assert next(iter(recorder.sources.values()))["status"] == "failed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_probe_failure_during_stop_is_recorded_and_capture_is_reaped(
|
|
tmp_path, fake_capture, monkeypatch
|
|
):
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
monkeypatch.setattr(
|
|
recording, "probe_video_segment", AsyncMock(side_effect=OSError("probe unavailable"))
|
|
)
|
|
await recorder.stop()
|
|
assert not recorder.active
|
|
assert all(process.returncode is not None for process in fake_capture.processes)
|
|
track = next(iter(recorder.sources.values()))
|
|
assert track["status"] == "failed"
|
|
assert "probe unavailable" in track["message"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_retries_process_cleanup_interrupted_in_watchdog(
|
|
tmp_path, fake_capture, monkeypatch
|
|
):
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
recorder._watchdog.cancel()
|
|
await asyncio.gather(recorder._watchdog, return_exceptions=True)
|
|
key = next(iter(recorder.active))
|
|
entered = asyncio.Event()
|
|
attempts = 0
|
|
terminate = recorder._terminate
|
|
|
|
async def interrupted_first(process):
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
entered.set()
|
|
await asyncio.Future()
|
|
await terminate(process)
|
|
|
|
monkeypatch.setattr(recorder, "_terminate", interrupted_first)
|
|
recorder._watchdog = asyncio.create_task(recorder._finish(key))
|
|
await entered.wait()
|
|
await recorder.stop()
|
|
assert attempts == 2
|
|
assert not recorder.active
|
|
assert all(process.returncode is not None for process in fake_capture.processes)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_preserves_clip_when_watchdog_is_cancelled_during_probe(
|
|
tmp_path, fake_capture, monkeypatch
|
|
):
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
recorder._watchdog.cancel()
|
|
await asyncio.gather(recorder._watchdog, return_exceptions=True)
|
|
key = next(iter(recorder.active))
|
|
entered = asyncio.Event()
|
|
attempts = 0
|
|
|
|
async def interrupted_probe(_output):
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
entered.set()
|
|
await asyncio.Future()
|
|
return {"duration": 2.5, "width": 1280, "height": 720}
|
|
|
|
monkeypatch.setattr(recording, "probe_video_segment", interrupted_probe)
|
|
recorder._watchdog = asyncio.create_task(recorder._finish(key))
|
|
await entered.wait()
|
|
await recorder.stop()
|
|
assert attempts == 2
|
|
assert not recorder.active
|
|
track = recorder.sources[key]
|
|
assert track["status"] == "ready"
|
|
assert track["segments"][0]["duration_ms"] == 2500
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resumed_recorder_appends_and_does_not_overwrite_previous_clips(
|
|
tmp_path, fake_capture
|
|
):
|
|
first = PluginRecordings(context(tmp_path))
|
|
await first.observe("source-1", "visual-test", source())
|
|
await first.stop()
|
|
original = next(iter(first.sources.values()))["segments"][0].copy()
|
|
original_path = tmp_path / original["file"]
|
|
original_path.write_bytes(b"immutable original clip")
|
|
second = PluginRecordings(context(tmp_path))
|
|
try:
|
|
await second.observe("source-1", "visual-test", source())
|
|
finally:
|
|
await second.stop()
|
|
segments = next(iter(second.sources.values()))["segments"]
|
|
assert len(segments) == 2
|
|
assert segments[0] == original
|
|
assert segments[1]["file"] != original["file"]
|
|
assert original_path.read_bytes() == b"immutable original clip"
|
|
manifest = json.loads(
|
|
(tmp_path / "evidence/recordings/recordings.json").read_text(encoding="utf-8")
|
|
)
|
|
assert len(manifest["sources"][0]["segments"]) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resized_window_seals_old_geometry_before_new_capture(tmp_path, fake_capture):
|
|
recorder = PluginRecordings(context(tmp_path, computer_scope="window"))
|
|
try:
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
await recorder.observe(
|
|
"source-2",
|
|
"visual-test",
|
|
source(window={"left": 40, "top": 50, "width": 1000, "height": 700}),
|
|
)
|
|
assert len(fake_capture.commands) == 2
|
|
assert fake_capture.processes[0].returncode is not None
|
|
assert len(recorder.active) == 1
|
|
finally:
|
|
await recorder.stop()
|
|
assert len(recorder.sources) == 2
|
|
assert all(track["status"] == "ready" for track in recorder.sources.values())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dead_worker_recovery_keeps_partial_clip_and_does_not_restart(
|
|
tmp_path, fake_capture, monkeypatch
|
|
):
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
await recorder.stop()
|
|
track = next(iter(recorder.sources.values()))
|
|
track.update(status="recording", owner_pid=987654321, owner_created_at=1)
|
|
track["segments"][0].update(duration=0, duration_ms=0)
|
|
recorder._save()
|
|
count = len(fake_capture.commands)
|
|
assert await recording.recover_interrupted_recordings(tmp_path)
|
|
recovered = PluginRecordings(context(tmp_path))
|
|
result = next(iter(recovered.sources.values()))
|
|
assert result["status"] == "failed"
|
|
assert result["segments"][0]["duration"] == 2.5
|
|
assert len(fake_capture.commands) == count
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_live_worker_recovery_never_finalizes_active_capture(tmp_path, fake_capture):
|
|
import os
|
|
import psutil
|
|
|
|
recorder = PluginRecordings(context(tmp_path))
|
|
await recorder.observe("source-1", "visual-test", source())
|
|
track = next(iter(recorder.sources.values()))
|
|
track.update(owner_pid=os.getpid(), owner_created_at=psutil.Process().create_time())
|
|
recorder._save()
|
|
try:
|
|
assert not await recording.recover_interrupted_recordings(tmp_path)
|
|
assert track["status"] == "recording"
|
|
finally:
|
|
await recorder.stop()
|