fix: Raise exception if multiple VFolders exist in decorator

Add changelog

Add test code for decorator

changes/3465.fix.md

@@ -0,0 +1 @@
+Raise exception if multiple VFolders exist in decorator

src/ai/backend/manager/api/vfolder.py

@@ -191,14 +191,13 @@ async def _wrapped(
             *args: P.args,
             **kwargs: P.kwargs,
         ) -> web.Response:
-            for row in folder_rows:
-                try:
-                    await check_vfolder_status(row, status)
-                    return await handler(request, row, *args, **kwargs)
-                except VFolderFilterStatusFailed:
-                    pass
-            # none of our candidates matched the status filter, so we should instead raise error here
-            raise VFolderFilterStatusFailed
+            if len(folder_rows) > 1:
+                raise TooManyVFoldersFound(folder_rows)
+            if len(folder_rows) == 0:
+                raise VFolderNotFound()
+            row = folder_rows[0]
+            await check_vfolder_status(row, status)
+            return await handler(request, row, *args, **kwargs)
 
         return _wrapped
 

src/ai/backend/manager/container_registry/base.py

@@ -186,10 +186,31 @@ async def commit_rescan_result(self) -> None:
                     progress_msg = f"Updated image - {parsed_img.canonical}/{image_identifier.architecture} ({update['config_digest']})"
                     log.info(progress_msg)
 
-                    if (reporter := progress_reporter.get()) is not None:
-                        await reporter.update(1, message=progress_msg)
+                    session.add(
+                        ImageRow(
+                            name=parsed_img.canonical,
+                            project=self.registry_info.project,
+                            registry=parsed_img.registry,
+                            registry_id=self.registry_info.id,
+                            image=join_non_empty(parsed_img.project, parsed_img.name, sep="/"),
+                            tag=parsed_img.tag,
+                            architecture=image_identifier.architecture,
+                            is_local=is_local,
+                            config_digest=update["config_digest"],
+                            size_bytes=update["size_bytes"],
+                            type=ImageType.COMPUTE,
+                            accelerators=update.get("accels"),
+                            labels=update["labels"],
+                            resources=update["resources"],
+                        )
+                    )
+                    progress_msg = f"Updated image - {parsed_img.canonical}/{image_identifier.architecture} ({update['config_digest']})"
+                    log.info(progress_msg)
+
+                if (reporter := progress_reporter.get()) is not None:
+                    await reporter.update(1, message=progress_msg)
 
-                await session.flush()
+            await session.flush()
 
     async def scan_single_ref(self, image: str) -> None:
         all_updates_token = all_updates.set({})

tests/manager/api/test_vfolder.py

@@ -1,11 +1,17 @@
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, Mock
 from uuid import UUID
 
 import pytest
+from aiohttp import web
 
 from ai.backend.manager.api import vfolder
-from ai.backend.manager.api.vfolder import with_vfolder_rows_resolved
-from ai.backend.manager.models.vfolder import VFolderPermissionSetAlias
+from ai.backend.manager.api.exceptions import TooManyVFoldersFound, VFolderNotFound
+from ai.backend.manager.api.vfolder import with_vfolder_rows_resolved, with_vfolder_status_checked
+from ai.backend.manager.models.vfolder import (
+    VFolderPermissionSetAlias,
+    VFolderRow,
+    VFolderStatusSet,
+)
 
 
 @pytest.mark.asyncio
@@ -27,3 +33,55 @@ async def dummy_handler(request, row):
     await dummy_handler(mock_request)
     call = mock_resolver.await_args_list[1]
     assert isinstance(call.args[2], str)
+
+
+@pytest.mark.parametrize(
+    "vfolder_status",
+    [
+        VFolderStatusSet.ALL,
+        VFolderStatusSet.READABLE,
+        VFolderStatusSet.MOUNTABLE,
+        VFolderStatusSet.UPDATABLE,
+        VFolderStatusSet.DELETABLE,
+        VFolderStatusSet.PURGABLE,
+        VFolderStatusSet.RECOVERABLE,
+        VFolderStatusSet.INACCESSIBLE,
+    ],
+)
+async def test_too_many_vfolders(vfolder_status):
+    @with_vfolder_status_checked(vfolder_status)
+    async def too_many_vfolders_handler(request, row: VFolderRow):
+        return AsyncMock(return_value=web.Response(text="no response"))
+
+    mock_entry = {
+        "id": "fake-vfolder-id",
+        "host": "fake-vfolder-host",
+        "user_email": "[email protected]",
+        "user": "fake-user",
+        "group_name": "fake-group",
+        "group": "fake-group-id",
+    }
+    with pytest.raises(TooManyVFoldersFound):
+        await too_many_vfolders_handler(Mock(), [mock_entry, mock_entry])
+
+
+@pytest.mark.parametrize(
+    "vfolder_status",
+    [
+        VFolderStatusSet.ALL,
+        VFolderStatusSet.READABLE,
+        VFolderStatusSet.MOUNTABLE,
+        VFolderStatusSet.UPDATABLE,
+        VFolderStatusSet.DELETABLE,
+        VFolderStatusSet.PURGABLE,
+        VFolderStatusSet.RECOVERABLE,
+        VFolderStatusSet.INACCESSIBLE,
+    ],
+)
+async def test_no_vfolders(vfolder_status):
+    @with_vfolder_status_checked(vfolder_status)
+    async def no_vfolders_handler(request, row: VFolderRow):
+        return AsyncMock(return_value=web.Response(text="no response"))
+
+    with pytest.raises(VFolderNotFound):
+        await no_vfolders_handler(Mock(), [])