fix(BA-517): Raise exception if multiple VFolders exist in 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
 

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(), [])