refactor

Author: sanderegg

services/storage/tests/fixtures/data_models.py

@@ -7,14 +7,14 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
 from random import choice, randint
-from typing import Any
+from typing import Any, cast
 
 import pytest
 import sqlalchemy as sa
 from faker import Faker
 from models_library.basic_types import SHA256Str
 from models_library.projects import ProjectAtDB, ProjectID
-from models_library.projects_nodes_io import NodeID, SimcoreS3FileID
+from models_library.projects_nodes_io import NodeID, SimcoreS3FileID, StorageFileID
 from models_library.users import UserID
 from pydantic import ByteSize, TypeAdapter
 from pytest_simcore.helpers.faker_factories import random_project, random_user
@@ -259,6 +259,39 @@ async def _creator(
     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, dict[str, Path | str]]],
+    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,
@@ -296,71 +329,66 @@ async def _creator(
     ) -> tuple[ProjectAtDB, dict[NodeID, dict[SimcoreS3FileID, dict[str, Path | str]]]]:
         assert len(file_sizes) == len(file_checksums)
         project = await create_project(name="random-project")
-        src_projects_list: dict[
+        node_to_files_mapping: dict[
             NodeID, dict[SimcoreS3FileID, dict[str, Path | str]]
         ] = {}
         upload_tasks: deque[Awaitable] = deque()
         for _node_index in range(num_nodes):
-            # NOTE: we put some more outputs in there to simulate a real case better
-            new_node_id = NodeID(f"{faker.uuid4()}")
+            # 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(
-                ProjectID(project["uuid"]),
-                new_node_id,
-                faker.file_name(),
-                Path("outputs/output3"),
+                project_id, node_id, output3_file_name, Path("outputs/output_3")
             )
-            src_node_id = await create_project_node(
+            created_node_id = await create_project_node(
                 ProjectID(project["uuid"]),
-                new_node_id,
+                node_id,
                 outputs={
                     "output_1": faker.pyint(),
                     "output_2": faker.pystr(),
                     "output_3": f"{output3_file_id}",
                 },
             )
-            assert src_node_id == new_node_id
-
-            # upload the output 3 and some random other files at the root of each node
-            src_projects_list[src_node_id] = {}
-            checksum: SHA256Str = choice(file_checksums)  # noqa: S311
-            src_file, _ = await upload_file(
-                file_size=choice(file_sizes),  # noqa: S311
-                file_name=Path(output3_file_id).name,
-                file_id=output3_file_id,
-                sha256_checksum=checksum,
-            )
-            src_projects_list[src_node_id][output3_file_id] = {
-                "path": src_file,
-                "sha256_checksum": checksum,
-            }
-
-            async def _upload_file_and_update_project(project, src_node_id):
-                src_file_name = faker.file_name()
-                src_file_uuid = create_simcore_file_id(
-                    ProjectID(project["uuid"]), src_node_id, src_file_name, None
+            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,
                 )
-                checksum: SHA256Str = choice(file_checksums)  # noqa: S311
-                src_file, _ = await upload_file(
-                    file_size=choice(file_sizes),  # noqa: S311
-                    file_name=src_file_name,
-                    file_id=src_file_uuid,
-                    sha256_checksum=checksum,
-                )
-                src_projects_list[src_node_id][src_file_uuid] = {
-                    "path": src_file,
-                    "sha256_checksum": checksum,
-                }
+            )
 
-            # add a few random files in the node storage
+            # add a few random files in the node workspace
             upload_tasks.extend(
                 [
-                    _upload_file_and_update_project(project, src_node_id)
+                    _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, src_projects_list
+        return project, node_to_files_mapping
 
     return _creator