cvat/tests/python/shared/utils/resource_import_export.py

import functools
import json
from abc import ABC, abstractmethod
from contextlib import ExitStack
from http import HTTPStatus
from time import sleep
from typing import Any, TypeVar

import pytest

T = TypeVar("T")

from shared.utils.config import get_method, post_method

FILENAME_TEMPLATE = "cvat/{}/{}.zip"
EXPORT_FORMAT = "CVAT for images 1.1"
IMPORT_FORMAT = "CVAT 1.1"


def _make_custom_resource_params(resource: str, obj: str, cloud_storage_id: int) -> dict[str, Any]:
    return {
        "filename": FILENAME_TEMPLATE.format(obj, resource),
        "location": "cloud_storage",
        "cloud_storage_id": cloud_storage_id,
    }


def _make_default_resource_params(resource: str, obj: str) -> dict[str, Any]:
    return {
        "filename": FILENAME_TEMPLATE.format(obj, resource),
    }


def _make_export_resource_params(
    resource: str, is_default: bool = True, **kwargs
) -> dict[str, Any]:
    func = _make_default_resource_params if is_default else _make_custom_resource_params
    params = func(resource, **kwargs)
    if resource != "backup":
        params["format"] = EXPORT_FORMAT
        params["save_images"] = resource == "dataset"

    return params


def _make_import_resource_params(
    resource: str, is_default: bool = True, **kwargs
) -> dict[str, Any]:
    func = _make_default_resource_params if is_default else _make_custom_resource_params
    params = func(resource, **kwargs)
    if resource != "backup":
        params["format"] = IMPORT_FORMAT
    return params


# FUTURE-TODO: reuse common logic from rest_api/utils
def _wait_request(
    user: str,
    request_id: str,
    *,
    sleep_interval: float = 0.1,
    number_of_checks: int = 100,
):
    for _ in range(number_of_checks):
        sleep(sleep_interval)
        response = get_method(user, f"requests/{request_id}")
        assert response.status_code == HTTPStatus.OK

        request_details = json.loads(response.content)
        status = request_details["status"]
        assert status in {"started", "queued", "finished", "failed"}
        if status in {"finished", "failed"}:
            return


class _CloudStorageResourceTest(ABC):
    @staticmethod
    @abstractmethod
    def _make_client():
        pass

    @pytest.fixture(autouse=True)
    def setup(self, admin_user: str):
        self.user = admin_user
        self.client = self._make_client()
        self.exit_stack = ExitStack()
        with self.exit_stack:
            yield

    def _ensure_file_created(self, func: T, storage: dict[str, Any]) -> T:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            filename = kwargs["filename"]
            bucket = storage["resource"]

            # check that file doesn't exist on the bucket
            assert not self.client.file_exists(bucket=bucket, filename=filename)

            func(*args, **kwargs)

            # check that file exists on the bucket
            assert self.client.file_exists(bucket=bucket, filename=filename)

        return wrapper

    def _export_resource_to_cloud_storage(
        self,
        obj_id: int,
        obj: str,
        resource: str,
        *,
        user: str,
        _expect_status: HTTPStatus = HTTPStatus.ACCEPTED,
        **kwargs,
    ):

        # initialize the export process
        response = post_method(
            user,
            f"{obj}/{obj_id}/{resource if resource != 'annotations' else 'dataset'}/export",
            data=None,
            **kwargs,
        )
        assert response.status_code == _expect_status

        if _expect_status == HTTPStatus.FORBIDDEN:
            return

        rq_id = json.loads(response.content).get("rq_id")
        assert rq_id, "The rq_id was not found in server request"

        _wait_request(user, rq_id)

    def _import_resource_from_cloud_storage(
        self, url: str, *, user: str, _expect_status: HTTPStatus = HTTPStatus.ACCEPTED, **kwargs
    ) -> None:
        response = post_method(user, url, data=None, **kwargs)
        status = response.status_code

        assert status == _expect_status, status
        if status == HTTPStatus.FORBIDDEN:
            return

        rq_id = response.json().get("rq_id")
        assert rq_id, "The rq_id parameter was not found in the server response"

        _wait_request(user, rq_id)

    def _import_annotations_from_cloud_storage(
        self,
        obj_id,
        obj,
        *,
        user,
        _expect_status: HTTPStatus = HTTPStatus.ACCEPTED,
        _check_uploaded: bool = True,
        **kwargs,
    ):
        url = f"{obj}/{obj_id}/annotations"
        self._import_resource_from_cloud_storage(
            url, user=user, _expect_status=_expect_status, **kwargs
        )

        if _expect_status == HTTPStatus.FORBIDDEN:
            return

        if _check_uploaded:
            response = get_method(user, url)
            assert response.status_code == HTTPStatus.OK

            annotations = response.json()

            assert len(annotations["shapes"])

    def _import_backup_from_cloud_storage(
        self, obj, *, user, _expect_status: HTTPStatus = HTTPStatus.ACCEPTED, **kwargs
    ):
        self._import_resource_from_cloud_storage(
            f"{obj}/backup", user=user, _expect_status=_expect_status, **kwargs
        )

    def _import_dataset_from_cloud_storage(
        self, obj_id, obj, *, user, _expect_status: HTTPStatus = HTTPStatus.ACCEPTED, **kwargs
    ):
        self._import_resource_from_cloud_storage(
            f"{obj}/{obj_id}/dataset", user=user, _expect_status=_expect_status, **kwargs
        )

    def _import_resource(self, cloud_storage: dict[str, Any], resource_type: str, *args, **kwargs):
        methods = {
            "annotations": self._import_annotations_from_cloud_storage,
            "dataset": self._import_dataset_from_cloud_storage,
            "backup": self._import_backup_from_cloud_storage,
        }

        org_id = cloud_storage["organization"]
        if org_id:
            kwargs.setdefault("org_id", org_id)

        kwargs.setdefault("user", self.user)

        return methods[resource_type](*args, **kwargs)

    def _export_resource(self, cloud_storage: dict[str, Any], *args, **kwargs):
        org_id = cloud_storage["organization"]
        if org_id:
            kwargs.setdefault("org_id", org_id)

        kwargs.setdefault("user", self.user)

        export_callback = self._ensure_file_created(
            self._export_resource_to_cloud_storage, storage=cloud_storage
        )
        export_callback(*args, **kwargs)

        self.exit_stack.callback(
            self.client.remove_file,
            bucket=cloud_storage["resource"],
            filename=kwargs["filename"],
        )