cvat-ai--cvat/tests/python/sdk/fixtures.py

# Copyright (C) CVAT.ai Corporation
#
# SPDX-License-Identifier: MIT

from pathlib import Path
from typing import Callable
from zipfile import ZipFile

import pytest
from cvat_sdk import Client
from cvat_sdk.api_client import models
from cvat_sdk.core.proxies.types import Location
from PIL import Image

from rest_api.utils import register_new_user
from shared.utils.config import BASE_URL, IMPORT_EXPORT_BUCKET_ID, USER_PASS, make_sdk_client
from shared.utils.helpers import generate_image_file

from .util import OrgResourceHierarchy, generate_coco_json


@pytest.fixture
def fxt_client(fxt_logger):
    logger, _ = fxt_logger

    client = Client(BASE_URL, logger=logger)
    api_client = client.api_client
    for k in api_client.configuration.logger:
        api_client.configuration.logger[k] = logger
    client.config.status_check_period = 0.01

    with client:
        yield client


@pytest.fixture
def fxt_image_file(tmp_path: Path):
    img_path = tmp_path / "img.png"
    with img_path.open("wb") as f:
        f.write(generate_image_file(filename=str(img_path), size=(5, 10)).getvalue())

    return img_path


@pytest.fixture
def fxt_coco_file(tmp_path: Path, fxt_image_file: Path):
    img_filename = fxt_image_file
    img_size = Image.open(img_filename).size
    ann_filename = tmp_path / "coco.json"
    generate_coco_json(ann_filename, img_info=(img_filename, *img_size))

    yield ann_filename


@pytest.fixture(scope="class")
def fxt_login(admin_user: str, restore_db_per_class):
    client = Client(BASE_URL)
    client.config.status_check_period = 0.01
    user = admin_user

    with client:
        client.login((user, USER_PASS))
        yield (client, user)


@pytest.fixture
def fxt_camvid_dataset(tmp_path: Path):
    img_path = tmp_path / "img.png"
    with img_path.open("wb") as f:
        f.write(generate_image_file(filename=str(img_path), size=(5, 10)).getvalue())

    annot_path = tmp_path / "annot.png"
    r, g, b = (127, 0, 0)
    annot = generate_image_file(
        filename=str(annot_path),
        size=(5, 10),
        color=(r, g, b),
    ).getvalue()
    with annot_path.open("wb") as f:
        f.write(annot)

    label_colors_path = tmp_path / "label_colors.txt"
    with open(label_colors_path, "w") as f:
        f.write(f"{r} {g} {b} car\n")

    dataset_img_path = "default/img.png"
    dataset_annot_path = "default/annot.png"
    default_txt_path = tmp_path / "default.txt"
    with open(default_txt_path, "w") as f:
        f.write(f"/{dataset_img_path} {dataset_annot_path}")

    dataset_path = tmp_path / "camvid_dataset.zip"
    with ZipFile(dataset_path, "x") as f:
        f.write(img_path, arcname=dataset_img_path)
        f.write(annot_path, arcname=dataset_annot_path)
        f.write(default_txt_path, arcname="default.txt")
        f.write(label_colors_path, arcname="label_colors.txt")

    yield dataset_path


@pytest.fixture
def fxt_coco_dataset(tmp_path: Path, fxt_image_file: Path, fxt_coco_file: Path):
    dataset_path = tmp_path / "coco_dataset.zip"
    with ZipFile(dataset_path, "x") as f:
        f.write(fxt_image_file, arcname="images/" + fxt_image_file.name)
        f.write(fxt_coco_file, arcname="annotations/instances_default.json")

    yield dataset_path


@pytest.fixture
def fxt_new_task(fxt_image_file: Path, fxt_login: tuple[Client, str]):
    client, _ = fxt_login
    task = client.tasks.create_from_data(
        spec={
            "name": "test_task",
            "labels": [{"name": "car"}, {"name": "person"}],
        },
        resources=[fxt_image_file],
        data_params={"image_quality": 80},
    )

    yield task


@pytest.fixture
def fxt_new_task_with_target_storage(fxt_image_file: Path, fxt_login: tuple[Client, str]):
    client, _ = fxt_login
    task = client.tasks.create_from_data(
        spec={
            "name": "test_task",
            "labels": [{"name": "car"}, {"name": "person"}],
            "target_storage": {
                "location": Location.CLOUD_STORAGE,
                "cloud_storage_id": IMPORT_EXPORT_BUCKET_ID,
            },
        },
        resources=[fxt_image_file],
        data_params={"image_quality": 80},
    )

    yield task


@pytest.fixture
def fxt_org_resource_hierarchy(
    fxt_image_file: Path,
) -> Callable[..., OrgResourceHierarchy]:
    def create(
        *, include_issue: bool = False, include_comment: bool = False
    ) -> OrgResourceHierarchy:
        owner = register_new_user("sdkowner")
        maintainer = register_new_user("sdkmaintainer")

        with make_sdk_client(owner["username"]) as owner_client:
            org = owner_client.organizations.create(models.OrganizationWriteRequest(slug="sdkorg"))
            owner_client.organization_slug = org.slug
            owner_client.api_client.invitations_api.create(
                models.InvitationWriteRequest(
                    role="maintainer",
                    email=maintainer["email"],
                ),
            )

            project = owner_client.projects.create(
                spec=models.ProjectWriteRequest(
                    name="org project",
                    labels=[models.PatchedLabelRequest(name="car")],
                )
            )
            task = owner_client.tasks.create_from_data(
                spec=models.TaskWriteRequest(name="org task", project_id=project.id),
                resources=[fxt_image_file],
                data_params={"image_quality": 80},
            )
            job = task.get_jobs()[0]

            issue_id = None
            comment_id = None

            if include_issue or include_comment:
                issue = owner_client.issues.create(
                    models.IssueWriteRequest(
                        frame=0,
                        position=[2.0, 4.0],
                        job=job.id,
                        message="hello",
                    )
                )
                issue_id = issue.id

                if include_comment:
                    comment = owner_client.comments.create(
                        models.CommentWriteRequest(issue.id, message="hi!")
                    )
                    comment_id = comment.id

        return OrgResourceHierarchy(
            owner_username=owner["username"],
            maintainer_username=maintainer["username"],
            org_slug=org.slug,
            project_id=project.id,
            task_id=task.id,
            job_id=job.id,
            issue_id=issue_id,
            comment_id=comment_id,
        )

    return create