pytest-simcore

Author: sanderegg

packages/pytest-simcore/src/pytest_simcore/helpers/storage_utils.py

@@ -1,10 +1,13 @@
 import logging
 import os
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, TypedDict
 
 import sqlalchemy as sa
+from faker import Faker
 from models_library.basic_types import SHA256Str
+from pydantic import ByteSize
 from simcore_postgres_database.storage_models import projects
 from sqlalchemy.ext.asyncio import AsyncEngine
 
@@ -33,3 +36,21 @@ async def get_updated_project(
 class FileIDDict(TypedDict):
     path: Path
     sha256_checksum: SHA256Str
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ProjectWithFilesParams:
+    num_nodes: int
+    allowed_file_sizes: tuple[ByteSize, ...]
+    workspace_files_count: int
+    allowed_file_checksums: tuple[SHA256Str, ...] = None  # type: ignore # NOTE: OK for testing
+
+    def __post_init__(self):
+        if self.allowed_file_checksums is None:
+            # generate some random checksums for the corresponding file sizes
+            faker = Faker()
+            checksums = tuple(faker.sha256() for _ in self.allowed_file_sizes)
+            object.__setattr__(self, "allowed_file_checksums", checksums)
+
+    def __repr__(self) -> str:
+        return f"ProjectWithFilesParams: #nodes={self.num_nodes}, file sizes={[_.human_readable() for _ in self.allowed_file_sizes]}"

packages/pytest-simcore/src/pytest_simcore/simcore_storage_data_models.py

@@ -2,29 +2,23 @@
 # pylint: disable=unused-argument
 # pylint: disable=unused-variable
 
-from collections import deque
 from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
-from pathlib import Path
-from random import choice, randint
-from typing import Any, cast
+from typing import Any
 
 import pytest
 import sqlalchemy as sa
 from faker import Faker
-from models_library.basic_types import SHA256Str
 from models_library.projects import ProjectID
-from models_library.projects_nodes_io import NodeID, SimcoreS3FileID, StorageFileID
+from models_library.projects_nodes_io import NodeID
 from models_library.users import UserID
-from pydantic import ByteSize, TypeAdapter
-from servicelib.utils import limited_gather
+from pydantic import TypeAdapter
 from simcore_postgres_database.models.project_to_groups import project_to_groups
 from simcore_postgres_database.storage_models import projects, users
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 
 from .helpers.faker_factories import random_project, random_user
-from .helpers.storage_utils import FileIDDict, get_updated_project
 
 
 @asynccontextmanager
@@ -257,134 +251,3 @@ async def _creator(
         return new_node_id
 
     return _creator
-
-
-async def _upload_file_and_update_project(
-    project_id: ProjectID,
-    node_id: NodeID,
-    *,
-    file_name: str | None,
-    file_id: StorageFileID | None,
-    file_sizes: tuple[ByteSize, ...],
-    file_checksums: tuple[SHA256Str, ...],
-    node_to_files_mapping: dict[NodeID, dict[SimcoreS3FileID, FileIDDict]],
-    upload_file: Callable[..., Awaitable[tuple[Path, SimcoreS3FileID]]],
-    create_simcore_file_id: Callable[
-        [ProjectID, NodeID, str, Path | None], SimcoreS3FileID
-    ],
-    faker: Faker,
-) -> None:
-    if file_name is None:
-        file_name = faker.file_name()
-        file_id = create_simcore_file_id(project_id, node_id, file_name, None)
-    checksum: SHA256Str = choice(file_checksums)  # noqa: S311
-    src_file, _ = await upload_file(
-        file_size=choice(file_sizes),  # noqa: S311
-        file_name=file_name,
-        file_id=file_id,
-        sha256_checksum=checksum,
-    )
-    assert file_name is not None
-    assert file_id is not None
-    node_to_files_mapping[node_id][file_id] = {
-        "path": src_file,
-        "sha256_checksum": checksum,
-    }
-
-
-@pytest.fixture
-async def random_project_with_files(
-    sqlalchemy_async_engine: AsyncEngine,
-    create_project: Callable[..., Awaitable[dict[str, Any]]],
-    create_project_node: Callable[..., Awaitable[NodeID]],
-    create_simcore_file_id: Callable[
-        [ProjectID, NodeID, str, Path | None], SimcoreS3FileID
-    ],
-    upload_file: Callable[..., Awaitable[tuple[Path, SimcoreS3FileID]]],
-    faker: Faker,
-) -> Callable[
-    [int, tuple[ByteSize, ...], tuple[SHA256Str, ...]],
-    Awaitable[tuple[dict[str, Any], dict[NodeID, dict[SimcoreS3FileID, FileIDDict]]]],
-]:
-    async def _creator(
-        num_nodes: int = 12,
-        file_sizes: tuple[ByteSize, ...] = (
-            TypeAdapter(ByteSize).validate_python("7Mib"),
-            TypeAdapter(ByteSize).validate_python("110Mib"),
-            TypeAdapter(ByteSize).validate_python("1Mib"),
-        ),
-        file_checksums: tuple[SHA256Str, ...] = (
-            TypeAdapter(SHA256Str).validate_python(
-                "311e2e130d83cfea9c3b7560699c221b0b7f9e5d58b02870bd52b695d8b4aabd"
-            ),
-            TypeAdapter(SHA256Str).validate_python(
-                "08e297db979d3c84f6b072c2a1e269e8aa04e82714ca7b295933a0c9c0f62b2e"
-            ),
-            TypeAdapter(SHA256Str).validate_python(
-                "488f3b57932803bbf644593bd46d95599b1d4da1d63bc020d7ebe6f1c255f7f3"
-            ),
-        ),
-    ) -> tuple[dict[str, Any], dict[NodeID, dict[SimcoreS3FileID, FileIDDict]]]:
-        assert len(file_sizes) == len(file_checksums)
-        project = await create_project(name="random-project")
-        node_to_files_mapping: dict[NodeID, dict[SimcoreS3FileID, FileIDDict]] = {}
-        upload_tasks: deque[Awaitable] = deque()
-        for _node_index in range(num_nodes):
-            # Create a node with outputs (files and others)
-            project_id = ProjectID(project["uuid"])
-            node_id = cast(NodeID, faker.uuid4(cast_to=None))
-            output3_file_name = faker.file_name()
-            output3_file_id = create_simcore_file_id(
-                project_id, node_id, output3_file_name, Path("outputs/output_3")
-            )
-            created_node_id = await create_project_node(
-                ProjectID(project["uuid"]),
-                node_id,
-                outputs={
-                    "output_1": faker.pyint(),
-                    "output_2": faker.pystr(),
-                    "output_3": f"{output3_file_id}",
-                },
-            )
-            assert created_node_id == node_id
-
-            node_to_files_mapping[created_node_id] = {}
-            upload_tasks.append(
-                _upload_file_and_update_project(
-                    project_id,
-                    node_id,
-                    file_name=output3_file_name,
-                    file_id=output3_file_id,
-                    file_sizes=file_sizes,
-                    file_checksums=file_checksums,
-                    upload_file=upload_file,
-                    create_simcore_file_id=create_simcore_file_id,
-                    faker=faker,
-                    node_to_files_mapping=node_to_files_mapping,
-                )
-            )
-
-            # add a few random files in the node workspace
-            upload_tasks.extend(
-                [
-                    _upload_file_and_update_project(
-                        project_id,
-                        node_id,
-                        file_name=None,
-                        file_id=None,
-                        file_sizes=file_sizes,
-                        file_checksums=file_checksums,
-                        upload_file=upload_file,
-                        create_simcore_file_id=create_simcore_file_id,
-                        faker=faker,
-                        node_to_files_mapping=node_to_files_mapping,
-                    )
-                    for _ in range(randint(0, 3))  # noqa: S311
-                ]
-            )
-        await limited_gather(*upload_tasks, limit=10)
-
-        project = await get_updated_project(sqlalchemy_async_engine, project["uuid"])
-        return project, node_to_files_mapping
-
-    return _creator

packages/pytest-simcore/src/pytest_simcore/simcore_storage_service.py

@@ -4,6 +4,7 @@
 import os
 from collections.abc import Callable, Iterable
 from copy import deepcopy
+from pathlib import Path
 
 import aiohttp
 import pytest
@@ -12,7 +13,6 @@
 from models_library.projects_nodes_io import NodeID, SimcoreS3FileID
 from pydantic import TypeAdapter
 from pytest_mock import MockerFixture
-from servicelib.minio_utils import ServiceRetryPolicyUponInitialization
 from yarl import URL
 
 from .helpers.docker import get_service_published_port
@@ -60,24 +60,34 @@ async def storage_service(
     return storage_endpoint
 
 
-# TODO: this can be used by ANY of the simcore services!
-@tenacity.retry(**ServiceRetryPolicyUponInitialization().kwargs)
+@tenacity.retry(
+    wait=tenacity.wait_fixed(1),
+    stop=tenacity.stop_after_delay(30),
+    reraise=True,
+)
 async def wait_till_storage_responsive(storage_endpoint: URL):
-    async with aiohttp.ClientSession() as session:
-        async with session.get(storage_endpoint.with_path("/v0/")) as resp:
-            assert resp.status == 200
-            data = await resp.json()
-            assert "data" in data
-            assert data["data"] is not None
+    async with (
+        aiohttp.ClientSession() as session,
+        session.get(storage_endpoint.with_path("/v0/")) as resp,
+    ):
+        assert resp.status == 200
+        data = await resp.json()
+        assert "data" in data
+        assert data["data"] is not None
 
 
 @pytest.fixture
 def create_simcore_file_id() -> Callable[[ProjectID, NodeID, str], SimcoreS3FileID]:
     def _creator(
-        project_id: ProjectID, node_id: NodeID, file_name: str
+        project_id: ProjectID,
+        node_id: NodeID,
+        file_name: str,
+        file_base_path: Path | None = None,
     ) -> SimcoreS3FileID:
-        return TypeAdapter(SimcoreS3FileID).validate_python(
-            f"{project_id}/{node_id}/{file_name}"
-        )
+        s3_file_name = file_name
+        if file_base_path:
+            s3_file_name = f"{file_base_path / file_name}"
+        clean_path = Path(f"{project_id}/{node_id}/{s3_file_name}")
+        return TypeAdapter(SimcoreS3FileID).validate_python(f"{clean_path}")
 
     return _creator