From 82adde42aa7727e22342e0b44132d5ef97beabea Mon Sep 17 00:00:00 2001 From: Maxim Zhiltsov Date: Thu, 8 Dec 2022 18:05:03 +0300 Subject: [PATCH] Refactor resource import export tests (#5429) Extracted some enhancements from https://github.com/opencv/cvat/pull/4819 - Extracted common s3 manipulations in tests - Refactored import/export tests to be more clear --- .../rest_api/test_resource_import_export.py | 258 ++++++++---------- tests/python/shared/utils/s3.py | 47 ++++ 2 files changed, 159 insertions(+), 146 deletions(-) create mode 100644 tests/python/shared/utils/s3.py diff --git a/tests/python/rest_api/test_resource_import_export.py b/tests/python/rest_api/test_resource_import_export.py index d1500a95..6e29a106 100644 --- a/tests/python/rest_api/test_resource_import_export.py +++ b/tests/python/rest_api/test_resource_import_export.py @@ -5,109 +5,131 @@ import functools import json +from contextlib import ExitStack from http import HTTPStatus +from typing import Any, Dict, TypeVar -import boto3 import pytest -from botocore.exceptions import ClientError -from shared.utils.config import ( - MINIO_ENDPOINT_URL, - MINIO_KEY, - MINIO_SECRET_KEY, - get_method, - post_method, -) +from shared.utils.config import get_method, post_method +from shared.utils.s3 import make_client + +T = TypeVar("T") FILENAME_TEMPLATE = "cvat/{}/{}.zip" FORMAT = "COCO 1.0" -def _use_custom_settings(obj, resource, cloud_storage_id): - return { +def _make_custom_resource_params(obj: str, resource: str, cloud_storage_id: int) -> Dict[str, Any]: + params = { "filename": FILENAME_TEMPLATE.format(obj, resource), - "use_default_location": False, "location": "cloud_storage", "cloud_storage_id": cloud_storage_id, - "format": FORMAT, + "use_default_location": False, } + if resource != "backup": + params["format"] = FORMAT + return params -def _use_default_settings(obj, resource): - return { +def _make_default_resource_params(obj: str, resource: str) -> Dict[str, Any]: + params = { "filename": FILENAME_TEMPLATE.format(obj, resource), "use_default_location": True, - "format": FORMAT, } + if resource != "backup": + params["format"] = FORMAT + return params -def define_client(): - s3 = boto3.resource( - "s3", - aws_access_key_id=MINIO_KEY, - aws_secret_access_key=MINIO_SECRET_KEY, - endpoint_url=MINIO_ENDPOINT_URL, - ) - return s3.meta.client +class _S3ResourceTest: + @pytest.fixture(autouse=True) + def setup(self, admin_user: str): + self.user = admin_user + self.s3_client = 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"] -def assert_file_does_not_exist(client, bucket, filename): - try: - client.head_object(Bucket=bucket, Key=filename) - raise AssertionError(f"File {filename} on bucket {bucket} already exists") - except ClientError: - pass + # check that file doesn't exist on the bucket + assert not self.s3_client.file_exists(bucket, filename) + func(*args, **kwargs) -def assert_file_exists(client, bucket, filename): - try: - client.head_object(Bucket=bucket, Key=filename) - except ClientError: - raise AssertionError(f"File {filename} on bucket {bucket} doesn't exist") + # check that file exists on the bucket + assert self.s3_client.file_exists(bucket, filename) + return wrapper -def assert_file_status(func): - @functools.wraps(func) - def wrapper(user, storage_conf, *args, **kwargs): - filename = kwargs["filename"] - bucket = storage_conf["resource"] - # get storage client - client = define_client() - # check that file doesn't exist on the bucket - assert_file_does_not_exist(client, bucket, filename) - func(user, storage_conf, *args, **kwargs) - # check that file exists on the bucket - assert_file_exists(client, bucket, filename) + def _export_resource_to_cloud_storage( + self, obj_id: int, obj: str, resource: str, *, user: str, **kwargs + ): + response = get_method(user, f"{obj}/{obj_id}/{resource}", **kwargs) + status = response.status_code - return wrapper + while status != HTTPStatus.OK: + assert status in (HTTPStatus.CREATED, HTTPStatus.ACCEPTED) + response = get_method(user, f"{obj}/{obj_id}/{resource}", action="download", **kwargs) + status = response.status_code + def _import_annotations_from_cloud_storage(self, obj_id, obj, *, user, **kwargs): + url = f"{obj}/{obj_id}/annotations" + response = post_method(user, url, data=None, **kwargs) + status = response.status_code -def remove_asset(bucket, filename): - client = define_client() - client.delete_object(Bucket=bucket, Key=filename) + while status != HTTPStatus.CREATED: + assert status == HTTPStatus.ACCEPTED + response = post_method(user, url, data=None, **kwargs) + status = response.status_code + def _import_backup_from_cloud_storage(self, obj_id, obj, *, user, **kwargs): + url = f"{obj}/backup" + response = post_method(user, url, data=None, **kwargs) + status = response.status_code -@assert_file_status -def _save_resource_to_cloud_storage(user, storage_conf, obj_id, obj, resource, **kwargs): - response = get_method(user, f"{obj}/{obj_id}/{resource}", **kwargs) - status = response.status_code + while status != HTTPStatus.CREATED: + assert status == HTTPStatus.ACCEPTED + data = json.loads(response.content.decode("utf8")) + response = post_method(user, url, data=data, **kwargs) + status = response.status_code - while status != HTTPStatus.OK: - assert status in (HTTPStatus.CREATED, HTTPStatus.ACCEPTED) - response = get_method(user, f"{obj}/{obj_id}/{resource}", action="download", **kwargs) + def _import_dataset_from_cloud_storage(self, obj_id, obj, *, user, **kwargs): + url = f"{obj}/{obj_id}/dataset" + response = post_method(user, url, data=None, **kwargs) status = response.status_code + while status != HTTPStatus.CREATED: + assert status == HTTPStatus.ACCEPTED + response = get_method(user, url, action="import_status") + status = response.status_code -def _idempotent_saving_resource_to_cloud_storage(*args, **kwargs): - _save_resource_to_cloud_storage(*args, **kwargs) - remove_asset(args[1]["resource"], kwargs["filename"]) + 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.s3_client.remove_file, + bucket=cloud_storage["resource"], + filename=kwargs["filename"], + ) -@pytest.mark.usefixtures("restore_db_per_class") -class TestSaveResource: - _USERNAME = "admin1" - _ORG = 2 +@pytest.mark.usefixtures("restore_db_per_class") +class TestExportResource(_S3ResourceTest): @pytest.mark.parametrize("cloud_storage_id", [3]) @pytest.mark.parametrize( "obj_id, obj, resource", @@ -126,13 +148,9 @@ class TestSaveResource: self, cloud_storage_id, obj_id, obj, resource, cloud_storages ): cloud_storage = cloud_storages[cloud_storage_id] - kwargs = _use_custom_settings(obj, resource, cloud_storage_id) - if resource == "backup": - kwargs.pop("format") + kwargs = _make_custom_resource_params(obj, resource, cloud_storage_id) - _idempotent_saving_resource_to_cloud_storage( - self._USERNAME, cloud_storage, obj_id, obj, resource, org_id=self._ORG, **kwargs - ) + self._export_resource(cloud_storage, obj_id, obj, resource, **kwargs) @pytest.mark.parametrize( "obj_id, obj, resource", @@ -168,56 +186,28 @@ class TestSaveResource: task_id = jobs[obj_id]["task_id"] cloud_storage_id = tasks[task_id]["target_storage"]["cloud_storage_id"] cloud_storage = cloud_storages[cloud_storage_id] + kwargs = _make_default_resource_params(obj, resource) - kwargs = _use_default_settings(obj, resource) - - if resource == "backup": - kwargs.pop("format") - - _idempotent_saving_resource_to_cloud_storage( - self._USERNAME, cloud_storage, obj_id, obj, resource, org_id=self._ORG, **kwargs - ) - - -def _import_annotations_from_cloud_storage(user, obj_id, obj, **kwargs): - url = f"{obj}/{obj_id}/annotations" - response = post_method(user, url, data=None, **kwargs) - status = response.status_code - - while status != HTTPStatus.CREATED: - assert status == HTTPStatus.ACCEPTED - response = post_method(user, url, data=None, **kwargs) - status = response.status_code - + self._export_resource(cloud_storage, obj_id, obj, resource, **kwargs) -def _import_backup_from_cloud_storage(user, obj_id, obj, **kwargs): - url = f"{obj}/backup" - response = post_method(user, url, data=None, **kwargs) - status = response.status_code - while status != HTTPStatus.CREATED: - assert status == HTTPStatus.ACCEPTED - data = json.loads(response.content.decode("utf8")) - response = post_method(user, url, data=data, **kwargs) - status = response.status_code - - -def _import_dataset_from_cloud_storage(user, obj_id, obj, **kwargs): - url = f"{obj}/{obj_id}/dataset" - response = post_method(user, url, data=None, **kwargs) - status = response.status_code +@pytest.mark.usefixtures("restore_db_per_function") +@pytest.mark.usefixtures("restore_cvat_data") +class TestImportResource(_S3ResourceTest): + 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, + } - while status != HTTPStatus.CREATED: - assert status == HTTPStatus.ACCEPTED - response = get_method(user, url, action="import_status") - status = response.status_code + org_id = cloud_storage["organization"] + if org_id: + kwargs.setdefault("org_id", org_id) + kwargs.setdefault("user", self.user) -@pytest.mark.usefixtures("restore_db_per_function") -@pytest.mark.usefixtures("restore_cvat_data") -class TestImportResource: - _USERNAME = "admin1" - _ORG = 2 + return methods[resource_type](*args, **kwargs) @pytest.mark.parametrize("cloud_storage_id", [3]) @pytest.mark.parametrize( @@ -234,26 +224,11 @@ class TestImportResource: self, cloud_storage_id, obj_id, obj, resource, cloud_storages ): cloud_storage = cloud_storages[cloud_storage_id] - kwargs = _use_custom_settings(obj, resource, cloud_storage_id) - export_kwargs = _use_custom_settings(obj, resource, cloud_storage_id) - - if resource == "backup": - kwargs.pop("format") - kwargs.pop("use_default_location") - export_kwargs.pop("format") + kwargs = _make_custom_resource_params(obj, resource, cloud_storage_id) + export_kwargs = _make_custom_resource_params(obj, resource, cloud_storage_id) + self._export_resource(cloud_storage, obj_id, obj, resource, **export_kwargs) - # export current resource to cloud storage - _save_resource_to_cloud_storage( - self._USERNAME, cloud_storage, obj_id, obj, resource, org_id=self._ORG, **export_kwargs - ) - - import_resource = { - "annotations": _import_annotations_from_cloud_storage, - "dataset": _import_dataset_from_cloud_storage, - "backup": _import_backup_from_cloud_storage, - } - import_resource[resource](self._USERNAME, obj_id, obj, org_id=self._ORG, **kwargs) - remove_asset(cloud_storage["resource"], kwargs["filename"]) + self._import_resource(cloud_storage, resource, obj_id, obj, **kwargs) @pytest.mark.parametrize( "obj_id, obj, resource", @@ -284,17 +259,8 @@ class TestImportResource: task_id = jobs[obj_id]["task_id"] cloud_storage_id = tasks[task_id]["source_storage"]["cloud_storage_id"] cloud_storage = cloud_storages[cloud_storage_id] - kwargs = _use_default_settings(obj, resource) - # export current resource to cloud storage - _save_resource_to_cloud_storage( - self._USERNAME, cloud_storage, obj_id, obj, resource, org_id=self._ORG, **kwargs - ) + kwargs = _make_default_resource_params(obj, resource) + self._export_resource(cloud_storage, obj_id, obj, resource, **kwargs) - import_resource = { - "annotations": _import_annotations_from_cloud_storage, - "dataset": _import_dataset_from_cloud_storage, - "backup": _import_backup_from_cloud_storage, - } - import_resource[resource](self._USERNAME, obj_id, obj, org_id=self._ORG, **kwargs) - remove_asset(cloud_storage["resource"], kwargs["filename"]) + self._import_resource(cloud_storage, resource, obj_id, obj, **kwargs) diff --git a/tests/python/shared/utils/s3.py b/tests/python/shared/utils/s3.py new file mode 100644 index 00000000..ee259081 --- /dev/null +++ b/tests/python/shared/utils/s3.py @@ -0,0 +1,47 @@ +# Copyright (C) 2022 CVAT.ai Corporation +# +# SPDX-License-Identifier: MIT + +import boto3 +from botocore.exceptions import ClientError + +from shared.utils.config import MINIO_ENDPOINT_URL, MINIO_KEY, MINIO_SECRET_KEY + + +class S3Client: + def __init__(self, endpoint_url: str, *, access_key: str, secret_key: str) -> None: + self.client = self._make_boto_client( + endpoint_url=endpoint_url, access_key=access_key, secret_key=secret_key + ) + + @staticmethod + def _make_boto_client(endpoint_url: str, *, access_key: str, secret_key: str): + s3 = boto3.resource( + "s3", + aws_access_key_id=access_key, + aws_secret_access_key=secret_key, + endpoint_url=endpoint_url, + ) + return s3.meta.client + + def create_file(self, bucket: str, filename: str, data: bytes = b""): + self.client.put_object(Body=data, Bucket=bucket, Key=filename) + + def remove_file(self, bucket: str, filename: str): + self.client.delete_object(Bucket=bucket, Key=filename) + + def file_exists(self, bucket: str, filename: str) -> bool: + try: + self.client.head_object(Bucket=bucket, Key=filename) + return True + except ClientError as e: + if e.response["Error"]["Code"] == "404": + return False + else: + raise + + +def make_client() -> S3Client: + return S3Client( + endpoint_url=MINIO_ENDPOINT_URL, access_key=MINIO_KEY, secret_key=MINIO_SECRET_KEY + )