core/tests/components/file_upload/test_init.py

"""Test the File Upload integration."""

from contextlib import contextmanager
from pathlib import Path
from random import getrandbits
from typing import Any
from unittest.mock import patch

import pytest

from homeassistant.components import file_upload
from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component

from tests.components.image_upload import TEST_IMAGE
from tests.typing import ClientSessionGenerator


@pytest.fixture
async def uploaded_file_dir(
    hass: HomeAssistant, hass_client: ClientSessionGenerator
) -> Path:
    """Test uploading and using a file."""
    assert await async_setup_component(hass, "file_upload", {})
    client = await hass_client()

    with (
        patch(
            # Patch temp dir name to avoid tests fail running in parallel
            "homeassistant.components.file_upload.TEMP_DIR_NAME",
            file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
        ),
        TEST_IMAGE.open("rb") as fp,
    ):
        res = await client.post("/api/file_upload", data={"file": fp})

    assert res.status == 200
    response = await res.json()

    file_dir = hass.data[file_upload.DOMAIN].file_dir(response["file_id"])
    assert file_dir.is_dir()
    return file_dir


async def test_using_file(hass: HomeAssistant, uploaded_file_dir) -> None:
    """Test uploading and using a file."""
    # Test we can use it
    with file_upload.process_uploaded_file(hass, uploaded_file_dir.name) as file_path:
        assert file_path.is_file()
        assert file_path.parent == uploaded_file_dir
        assert file_path.read_bytes() == TEST_IMAGE.read_bytes()

    # Test it's removed
    assert not uploaded_file_dir.exists()


async def test_removing_file(
    hass: HomeAssistant, hass_client: ClientSessionGenerator, uploaded_file_dir
) -> None:
    """Test uploading and using a file."""
    client = await hass_client()

    response = await client.delete(
        "/api/file_upload", json={"file_id": uploaded_file_dir.name}
    )
    assert response.status == 200

    # Test it's removed
    assert not uploaded_file_dir.exists()


async def test_removed_on_stop(
    hass: HomeAssistant, hass_client: ClientSessionGenerator, uploaded_file_dir
) -> None:
    """Test uploading and using a file."""
    await hass.async_stop()

    # Test it's removed
    assert not uploaded_file_dir.exists()


async def test_upload_large_file(
    hass: HomeAssistant, hass_client: ClientSessionGenerator, large_file_io
) -> None:
    """Test uploading large file."""
    assert await async_setup_component(hass, "file_upload", {})
    client = await hass_client()

    with (
        patch(
            # Patch temp dir name to avoid tests fail running in parallel
            "homeassistant.components.file_upload.TEMP_DIR_NAME",
            file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
        ),
        patch(
            # Patch one megabyte to 50 bytes to prevent having to use big files in tests
            "homeassistant.components.file_upload.ONE_MEGABYTE",
            50,
        ),
    ):
        res = await client.post("/api/file_upload", data={"file": large_file_io})

    assert res.status == 200
    response = await res.json()

    file_dir = hass.data[file_upload.DOMAIN].file_dir(response["file_id"])
    assert file_dir.is_dir()

    large_file_io.seek(0)
    with file_upload.process_uploaded_file(hass, file_dir.name) as file_path:
        assert file_path.is_file()
        assert file_path.parent == file_dir
        assert file_path.read_bytes() == large_file_io.read().encode("utf-8")


async def test_upload_with_wrong_key_fails(
    hass: HomeAssistant, hass_client: ClientSessionGenerator, large_file_io
) -> None:
    """Test uploading fails."""
    assert await async_setup_component(hass, "file_upload", {})
    client = await hass_client()

    with patch(
        # Patch temp dir name to avoid tests fail running in parallel
        "homeassistant.components.file_upload.TEMP_DIR_NAME",
        file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
    ):
        res = await client.post("/api/file_upload", data={"wrong_key": large_file_io})

    assert res.status == 400


async def test_upload_large_file_fails(
    hass: HomeAssistant, hass_client: ClientSessionGenerator, large_file_io
) -> None:
    """Test uploading large file."""
    assert await async_setup_component(hass, "file_upload", {})
    client = await hass_client()

    @contextmanager
    def _mock_open(*args, **kwargs):
        yield MockPathOpen()

    class MockPathOpen:
        def __init__(self, *args: Any, **kwargs: Any) -> None:
            pass

        def write(self, data: bytes) -> None:
            raise OSError("Boom")

    with (
        patch(
            # Patch temp dir name to avoid tests fail running in parallel
            "homeassistant.components.file_upload.TEMP_DIR_NAME",
            file_upload.TEMP_DIR_NAME + f"-{getrandbits(10):03x}",
        ),
        patch(
            # Patch one megabyte to 50 bytes to prevent having to use big files in tests
            "homeassistant.components.file_upload.ONE_MEGABYTE",
            50,
        ),
        patch(
            "homeassistant.components.file_upload.Path.open", return_value=_mock_open()
        ),
    ):
        res = await client.post("/api/file_upload", data={"file": large_file_io})

    assert res.status == 500

    response = await res.content.read()

    assert b"Boom" in response