FrameTour-RenderWorker/tests/unit/test_base_handler_parallel_transfer.py

# -*- coding: utf-8 -*-

import os
from contextlib import contextmanager
from types import SimpleNamespace

import pytest

from domain.config import WorkerConfig
from domain.result import TaskResult
from domain.task import TaskType
from handlers.base import BaseHandler


class _DummyApiClient:
    pass


class _DummyHandler(BaseHandler):
    def handle(self, task):
        return TaskResult.ok({})

    def get_supported_type(self):
        return TaskType.RENDER_SEGMENT_TS


def _create_handler(tmp_path):
    config = WorkerConfig(
        api_endpoint='http://127.0.0.1:18084/api',
        access_key='TEST_ACCESS_KEY',
        worker_id='test-worker',
        temp_dir=str(tmp_path),
        cache_enabled=False,
        cache_dir=str(tmp_path / 'cache')
    )
    return _DummyHandler(config, _DummyApiClient())


def test_download_files_parallel_collects_success_and_failure(tmp_path, monkeypatch):
    handler = _create_handler(tmp_path)
    handler.task_download_concurrency = 3
    captured_calls = []

    def _fake_download(url, dest, timeout=None, use_cache=True):
        captured_calls.append((url, dest, timeout, use_cache))
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        with open(dest, 'wb') as file_obj:
            file_obj.write(b'1')
        return not url.endswith('/fail')

    monkeypatch.setattr(handler, 'download_file', _fake_download)
    results = handler.download_files_parallel(
        [
            {'key': 'first', 'url': 'https://example.com/first', 'dest': str(tmp_path / 'first.bin')},
            {'key': 'second', 'url': 'https://example.com/fail', 'dest': str(tmp_path / 'second.bin')},
            {'key': 'third', 'url': 'https://example.com/third', 'dest': str(tmp_path / 'third.bin'), 'use_cache': False},
        ],
        timeout=15,
    )

    assert len(captured_calls) == 3
    assert results['first']['success'] is True
    assert results['second']['success'] is False
    assert results['third']['success'] is True
    assert any(call_item[3] is False for call_item in captured_calls)


def test_download_files_parallel_rejects_duplicate_key(tmp_path):
    handler = _create_handler(tmp_path)
    with pytest.raises(ValueError, match='Duplicate download job key'):
        handler.download_files_parallel(
            [
                {'key': 'dup', 'url': 'https://example.com/1', 'dest': str(tmp_path / '1.bin')},
                {'key': 'dup', 'url': 'https://example.com/2', 'dest': str(tmp_path / '2.bin')},
            ]
        )


def test_upload_files_parallel_collects_urls(tmp_path, monkeypatch):
    handler = _create_handler(tmp_path)
    handler.task_upload_concurrency = 2

    def _fake_upload(task_id, file_type, file_path, file_name=None):
        if file_type == 'video':
            return f'https://cdn.example.com/{task_id}/{file_name or "video.mp4"}'
        return None

    monkeypatch.setattr(handler, 'upload_file', _fake_upload)
    results = handler.upload_files_parallel(
        [
            {
                'key': 'video_output',
                'task_id': 'task-1',
                'file_type': 'video',
                'file_path': str(tmp_path / 'video.mp4'),
                'file_name': 'output.mp4',
            },
            {
                'key': 'audio_output',
                'task_id': 'task-1',
                'file_type': 'audio',
                'file_path': str(tmp_path / 'audio.aac'),
            },
        ]
    )

    assert results['video_output']['success'] is True
    assert results['video_output']['url'] == 'https://cdn.example.com/task-1/output.mp4'
    assert results['audio_output']['success'] is False
    assert results['audio_output']['url'] is None


def test_download_file_sets_lock_wait_ms_span_attribute(tmp_path, monkeypatch):
    handler = _create_handler(tmp_path)
    destination = tmp_path / "download.bin"

    class _FakeSpan:
        def __init__(self):
            self.attributes = {}

        def set_attribute(self, key, value):
            self.attributes[key] = value

    fake_span = _FakeSpan()

    @contextmanager
    def _fake_start_span(name, kind=None, attributes=None):
        if attributes:
            fake_span.attributes.update(attributes)
        yield fake_span

    def _fake_get_or_download_with_metrics(url, dest, timeout=300, max_retries=5):
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        with open(dest, 'wb') as file_obj:
            file_obj.write(b'abc')
        return True, {"lock_wait_ms": 1234, "lock_acquired": True, "cache_path_used": "cache"}

    monkeypatch.setattr("handlers.base.start_span", _fake_start_span)
    monkeypatch.setattr(
        handler.material_cache,
        "get_or_download_with_metrics",
        _fake_get_or_download_with_metrics
    )

    assert handler.download_file("https://example.com/file.bin", str(destination), timeout=1, use_cache=True)
    assert fake_span.attributes["render.file.lock_wait_ms"] == 1234
    assert fake_span.attributes["render.file.lock_acquired"] is True
    assert fake_span.attributes["render.file.cache_path_used"] == "cache"


def test_download_file_without_cache_sets_lock_wait_ms_zero(tmp_path, monkeypatch):
    handler = _create_handler(tmp_path)
    destination = tmp_path / "download-no-cache.bin"

    class _FakeSpan:
        def __init__(self):
            self.attributes = {}

        def set_attribute(self, key, value):
            self.attributes[key] = value

    fake_span = _FakeSpan()

    @contextmanager
    def _fake_start_span(name, kind=None, attributes=None):
        if attributes:
            fake_span.attributes.update(attributes)
        yield fake_span

    def _fake_storage_download(url, dest, timeout=30):
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        with open(dest, 'wb') as file_obj:
            file_obj.write(b'def')
        return True

    monkeypatch.setattr("handlers.base.start_span", _fake_start_span)
    monkeypatch.setattr("handlers.base.storage.download_file", _fake_storage_download)

    assert handler.download_file("https://example.com/file.bin", str(destination), timeout=1, use_cache=False)
    assert fake_span.attributes["render.file.lock_wait_ms"] == 0
    assert fake_span.attributes["render.file.lock_acquired"] is False
    assert fake_span.attributes["render.file.cache_path_used"] == "direct"


def test_upload_file_sets_detailed_span_attributes(tmp_path, monkeypatch):
    handler = _create_handler(tmp_path)
    source_path = tmp_path / "upload.mp4"
    source_path.write_bytes(b"abc123")
    fake_span_attributes = {}

    class _FakeSpan:
        def set_attribute(self, key, value):
            fake_span_attributes[key] = value

    @contextmanager
    def _fake_start_span(name, kind=None, attributes=None):
        if attributes:
            fake_span_attributes.update(attributes)
        yield _FakeSpan()

    handler.api_client = SimpleNamespace(
        get_upload_url=lambda *args, **kwargs: {
            "uploadUrl": "https://example.com/upload",
            "accessUrl": "https://cdn.example.com/output.mp4",
        }
    )

    monkeypatch.setattr("handlers.base.start_span", _fake_start_span)
    monkeypatch.setattr(
        "handlers.base.storage.upload_file_with_metrics",
        lambda *args, **kwargs: (
            True,
            {
                "upload_method": "http",
                "http_attempts": 2,
                "http_retry_count": 1,
                "http_status_code": 200,
                "http_replace_applied": True,
                "content_type": "video/mp4",
                "error_type": "",
                "rclone_attempted": False,
                "rclone_succeeded": False,
                "rclone_fallback_http": False,
            },
        )
    )
    monkeypatch.setattr(handler.material_cache, "add_to_cache", lambda *args, **kwargs: True)

    access_url = handler.upload_file("task-1", "video", str(source_path), "output.mp4")
    assert access_url == "https://cdn.example.com/output.mp4"
    assert fake_span_attributes["render.file.upload_success"] is True
    assert fake_span_attributes["render.file.upload_method"] == "http"
    assert fake_span_attributes["render.file.http_attempts"] == 2
    assert fake_span_attributes["render.file.http_retry_count"] == 1
    assert fake_span_attributes["render.file.http_status_code"] == 200
    assert fake_span_attributes["render.file.http_replace_applied"] is True
    assert fake_span_attributes["render.file.content_type"] == "video/mp4"
    assert fake_span_attributes["render.file.cache_write_back"] == "success"