wrolpi/modules/videos/conftest.py

import json
import pathlib
import shutil
from http import HTTPStatus
from typing import List
from uuid import uuid4

import mock
import pytest
from PIL import Image

from modules.videos.downloader import VideoDownloader, ChannelDownloader
from modules.videos.lib import set_test_channels_config, set_test_downloader_config, format_videos_destination
from modules.videos.models import Channel, Video
from wrolpi.api_utils import api_app
from wrolpi.cmd import CommandResult
from wrolpi.downloader import DownloadFrequency, DownloadManager
from wrolpi.files.models import FileGroup, Directory
from wrolpi.vars import PROJECT_DIR


@pytest.fixture
def simple_channel(test_session, test_directory) -> Channel:
    """Get a Channel with the minimum properties.  This Channel has no download!"""
    from wrolpi.collections import Collection

    # Create Collection first
    collection = Collection(
        name='Simple Channel',
        kind='channel',
        directory=test_directory,
    )
    test_session.add(collection)
    test_session.flush([collection])

    # Create Channel linked to Collection
    channel = Channel(
        collection_id=collection.id,
        url='https://example.com/channel1',
    )

    test_session.add(channel)
    test_session.commit()
    return channel


@pytest.fixture
def channel_factory(test_session, test_directory):
    """Create a random Channel with a directory, and Download."""
    from wrolpi.tags import Tag
    from wrolpi.collections import Collection

    def factory(source_id: str = None, download_frequency: DownloadFrequency = None, url: str = None, name: str = None,
                directory: pathlib.Path = None, tag_name: str = None):
        name = name or str(uuid4())
        tag = Tag.find_by_name(test_session, tag_name) if tag_name else None
        tag_name = tag.name if tag else None
        directory = directory or format_videos_destination(name, tag_name, url or f'https://example.com/{name}')
        directory.mkdir(exist_ok=True, parents=True)

        # Create Collection first
        collection = Collection(
            name=name,
            kind='channel',
            directory=directory,
            tag_id=tag.id if tag else None,
        )
        test_session.add(collection)
        test_session.flush([collection])

        # Create Channel linked to Collection
        channel = Channel(
            collection_id=collection.id,
            url=url or f'https://example.com/{name}',
            source_id=source_id,
        )
        test_session.add(channel)
        test_session.flush([channel])
        test_session.add(Directory(path=channel.directory, name=channel.directory.name))
        assert channel.id and channel.url
        if download_frequency:
            download = channel.get_or_create_download(test_session, channel.url, download_frequency)
            assert download.url == channel.url
            assert channel.downloads
        test_session.commit()
        return channel

    return factory


@pytest.fixture
def download_channel(test_session, test_directory, video_download_manager) -> Channel:
    """Get a test Channel that has a download frequency."""
    from wrolpi.collections import Collection

    # Create Collection first
    collection = Collection(
        name='Download Channel',
        kind='channel',
        directory=test_directory,
    )
    test_session.add(collection)
    test_session.flush([collection])

    # Create Channel linked to Collection
    channel = Channel(
        collection_id=collection.id,
        url='https://example.com/channel1',
        source_id='channel1',
    )
    test_session.add(channel)
    test_session.flush([channel])
    assert channel and channel.id and channel.url
    download = channel.get_or_create_download(test_session, channel.url, 60)
    assert download.url == channel.url
    download.frequency = DownloadFrequency.weekly
    test_session.commit()
    return channel


@pytest.fixture
def simple_video(test_session, test_directory, simple_channel, video_file) -> Video:
    """A Video with a video file, the channel is `simple_channel`."""
    video_path = test_directory / 'simple_video.mp4'
    video_file.rename(video_path)
    video = Video.from_paths(test_session, video_path)
    video.channel = simple_channel
    test_session.commit()
    video = Video.find_by_id(test_session, video.id)
    return video


@pytest.fixture
def video_factory(test_session, test_directory):
    """Creates Videos for testing."""

    def factory(channel_id: int = None, title: str = None, upload_date=None, with_video_file=None,
                with_info_json: dict = None, with_poster_ext: str = None, with_caption_file: bool = False,
                source_id: str = None, tag_names: List[str] = None) -> Video:
        title = title or str(uuid4())
        tag_names = tag_names or list()

        if with_video_file and isinstance(with_video_file, (pathlib.Path, str)):
            # Put the video exactly where specified for the test.
            path = pathlib.Path(with_video_file) if not isinstance(with_video_file, pathlib.Path) else with_video_file
        elif channel_id:
            # Put the video in its Channel's directory.
            channel = Channel.find_by_id(test_session, channel_id)
            path = (channel.directory or test_directory) / f'{title}.mp4'
            path.parent.mkdir(exist_ok=True, parents=True)
        else:
            # Put any video not in a Channel and without a specified file in the NO CHANNEL directory.
            (test_directory / 'videos/NO CHANNEL').mkdir(exist_ok=True, parents=True)
            path = test_directory / f'videos/NO CHANNEL/{title}.mp4'

        assert str(path).startswith(str(test_directory)), 'Video must be in the test directory'

        # Create a real video file for mimetype.
        shutil.copy(PROJECT_DIR / 'test/big_buck_bunny_720p_1mb.mp4', path)

        info_json_path = None
        if with_info_json:
            # Use the provided object as the json.
            with_info_json = {'duration': 5, 'epoch': 12345678} if with_info_json is True else with_info_json
            info_json_path = path.with_suffix('.info.json')
            info_json_path.write_text(json.dumps(with_info_json))

        poster_path = None
        if with_poster_ext:
            poster_path = path.with_suffix(f'.{with_poster_ext.lstrip(".")}')
            Image.new('RGB', (25, 25)).save(poster_path)

        caption_path = None
        if with_caption_file:
            caption_path = path.with_suffix('.en.vtt')
            shutil.copy(PROJECT_DIR / 'test/example1.en.vtt', caption_path)

        paths = (path, info_json_path, poster_path, caption_path)
        paths = list(filter(None, paths))

        video = Video.from_paths(test_session, *paths)
        video.channel_id = channel_id
        video.source_id = source_id or title
        video.file_group.published_datetime = upload_date
        video.validate(test_session)

        for tag_name in tag_names:
            video.add_tag(test_session, tag_name)

        return video

    return factory


@pytest.fixture
def audio_factory(test_session, test_directory):
    """Creates audio-only Video records for testing."""

    def factory(channel_id: int = None, title: str = None, upload_date=None,
                with_info_json: dict = None, with_poster_ext: str = None,
                source_id: str = None, tag_names: List[str] = None) -> Video:
        title = title or str(uuid4())
        tag_names = tag_names or list()

        if channel_id:
            channel = Channel.find_by_id(test_session, channel_id)
            path = (channel.directory or test_directory) / f'{title}.mp3'
            path.parent.mkdir(exist_ok=True, parents=True)
        else:
            (test_directory / 'videos/NO CHANNEL').mkdir(exist_ok=True, parents=True)
            path = test_directory / f'videos/NO CHANNEL/{title}.mp3'

        assert str(path).startswith(str(test_directory)), 'Audio must be in the test directory'

        shutil.copy(PROJECT_DIR / 'test/big_buck_bunny.mp3', path)

        info_json_path = None
        if with_info_json:
            with_info_json = {'duration': 5, 'epoch': 12345678} if with_info_json is True else with_info_json
            info_json_path = path.with_suffix('.info.json')
            info_json_path.write_text(json.dumps(with_info_json))

        poster_path = None
        if with_poster_ext:
            poster_path = path.with_suffix(f'.{with_poster_ext.lstrip(".")}')
            Image.new('RGB', (25, 25)).save(poster_path)

        paths = list(filter(None, (path, info_json_path, poster_path)))

        video = Video.from_paths(test_session, *paths)
        video.channel_id = channel_id
        video.source_id = source_id or title
        video.file_group.published_datetime = upload_date
        video.validate(test_session)

        for tag_name in tag_names:
            video.add_tag(test_session, tag_name)

        return video

    return factory


@pytest.fixture
def video_download_manager(test_download_manager) -> DownloadManager:
    """Get a DownloadManager ready to download Videos and Channels."""
    channel_downloader = ChannelDownloader()
    video_downloader = VideoDownloader()

    test_download_manager.register_downloader(channel_downloader)
    test_download_manager.register_downloader(video_downloader)

    with mock.patch('modules.videos.downloader.YDL.extract_info') as mock_extract_info:  # prevent real downloads
        mock_extract_info.side_effect = Exception('You must patch modules.videos.downloader.YDL.extract_info!')
        yield test_download_manager


@pytest.fixture
def test_channels_config(test_directory, async_client):
    (test_directory / 'config').mkdir(exist_ok=True)
    config_path = test_directory / 'config/channels.yaml'
    with set_test_channels_config() as config:
        config.initialize(api_app.shared_ctx.channels_config)
        yield config_path


@pytest.fixture
def test_videos_downloader_config(test_directory):
    (test_directory / 'config').mkdir(exist_ok=True)
    config_path = test_directory / 'config/videos_downloader.yaml'
    set_test_downloader_config(True)
    yield config_path
    set_test_downloader_config(False)


@pytest.fixture
def mock_video_extract_info():
    with mock.patch('modules.videos.downloader.extract_info') as mock_extract_info:
        # Add some simple data so the function can be called.
        mock_extract_info.return_value = dict(
            entries=[],
            uploader='mock_video_extract_info',
            id='mock_video_extract_info',
        )
        yield mock_extract_info


@pytest.fixture
def mock_video_prepare_filename():
    with mock.patch('modules.videos.downloader.prepare_filename') as mock_prepare_filename:
        yield mock_prepare_filename


@pytest.fixture
def mock_video_process_runner():
    with mock.patch('modules.videos.downloader.VideoDownloader.process_runner') as mock_process_runner:
        mock_process_runner.return_value = \
            CommandResult(return_code=0, cancelled=False, stdout=b'', stderr=b'', elapsed=0)
        yield mock_process_runner


@pytest.fixture
def video_with_search_factory(test_session, test_directory, video_file_factory):
    """A factory that creates a Video record with an associated video File record."""

    def video_with_search(path: str = None, title: str = None, b_text: str = None, c_text: str = None,
                          d_text: str = None, mimetype: str = 'video/mp4'):
        video_file_group = FileGroup.from_paths(test_session, video_file_factory(path))
        video_file_group.a_text = title
        video_file_group.b_text = b_text
        video_file_group.c_text = c_text
        video_file_group.d_text = d_text
        video_file_group.mimetype = mimetype
        video_file_group.model = 'video'
        video = Video(file_group=video_file_group)
        test_session.add(video)
        return video

    return video_with_search


@pytest.fixture
def assert_video_search(async_client):
    """A fixture which performs a video search request against the API.

    If assert_* params are passed, the response is checked."""

    async def _search_videos(
            assert_total: int = None,
            assert_ids: List[int] = None,
            assert_paths: List[str] = None,
            search_str: str = None,
            tag_names: List[str] = None,
            offset: int = None,
            limit: int = None,
            order_by: str = None,
            channel_id: int = None,
            censored: bool = None,
            deep: bool = None,
    ):
        content = dict()
        if search_str is not None:
            content['search_str'] = search_str
        if tag_names is not None:
            content['tag_names'] = tag_names
        if offset is not None:
            content['offset'] = offset
        if limit is not None:
            content['limit'] = limit
        if order_by is not None:
            content['order_by'] = order_by
        if channel_id is not None:
            content['channel_id'] = channel_id
        if censored is not None:
            content['censored'] = censored
        if deep is not None:
            content['deep'] = deep

        request, response = await async_client.post('/api/videos/search', content=json.dumps(content))

        assert response.json
        assert response.status_code == HTTPStatus.OK

        if assert_ids is not None:
            assert 'file_groups' in response.json, 'No video file_groups in response'
            assert [i['id'] for i in response.json['file_groups']] == assert_ids, 'Video IDs do not match'

        if assert_paths is not None:
            assert 'file_groups' in response.json, 'No video file_groups in response'
            assert [i['primary_path'] for i in response.json['file_groups']] == assert_paths, 'Video paths do not match'

        if assert_total and not (count := response.json['totals']['file_groups']) == int(assert_total):
            raise AssertionError(f'Video count does not match: {count} != {assert_total}')

        return request, response

    return _search_videos