UN-4398 [MISC] Ask the active auth service for roles in backend tests

backend/account_v2/test_organization_list.py

@@ -9,6 +9,7 @@
 from rest_framework.request import Request
 from rest_framework.test import APIRequestFactory, force_authenticate
 from tenant_account_v2.models import OrganizationMember
+from utils.tests.auth_seam import builtin_auth_service, patch_auth_service
 from utils.user_context import UserContext
 
 from account_v2 import urls, views
@@ -49,21 +50,11 @@ def setUp(self) -> None:
         UserContext.set_organization_identifier(org_a.organization_id)
         self.addCleanup(UserContext.set_organization_identifier, None)
 
-        for target, kwargs in [
-            (
-                "account_v2.authentication_controller."
-                "AuthenticationPluginRegistry.is_plugin_available",
-                {"return_value": False},
-            ),
-            (
-                "account_v2.authentication_controller."
-                "CacheService.set_user_organizations",
-                {},
-            ),
-        ]:
-            patcher = patch(target, **kwargs)
-            patcher.start()
-            self.addCleanup(patcher.stop)
+        patcher = patch(
+            "account_v2.authentication_controller.CacheService.set_user_organizations"
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
 
         self.controller = AuthenticationController()
         self.listed = [
@@ -195,9 +186,8 @@ def test_view_returns_the_old_payload_plus_details(self) -> None:
         force_authenticate(raw, user=self.caller)
         # org-x has no data on this deployment, as for another region's org.
         unknown = OrganizationData(id="org-x", display_name="X", name="org-x")
-        with patch(
-            "account_v2.authentication_controller.AuthenticationService.user_organizations",
-            return_value=[*self.listed, unknown],
+        with patch_auth_service(
+            "user_organizations", return_value=[*self.listed, unknown]
         ):
             response = views.get_organizations(raw)
         self.assertEqual(response.status_code, 200)
@@ -288,23 +278,16 @@ def setUp(self) -> None:
             for o in (org_a, org_b)
         ]
         self.list_organizations = MagicMock(return_value=listed)
+        # The built-in session auth throughout: a plugin's middleware would 401
+        # a session it did not seed. Test sessions have no plugin hook yet.
+        self.enterContext(builtin_auth_service())
         for target, value in [
             (
                 "account_v2.authentication_controller."
                 "AuthenticationService.user_organizations",
                 self.list_organizations,
             ),
-            # OSS auth throughout, and no write to the shared cache.
-            (
-                "account_v2.custom_auth_middleware."
-                "AuthenticationPluginRegistry.is_plugin_available",
-                MagicMock(return_value=False),
-            ),
-            (
-                "account_v2.authentication_controller."
-                "AuthenticationPluginRegistry.is_plugin_available",
-                MagicMock(return_value=False),
-            ),
+            # No write to the shared cache.
             (
                 "account_v2.authentication_controller."
                 "CacheService.set_user_organizations",

backend/account_v2/test_product_preference.py

@@ -27,8 +27,8 @@
 from django.db import IntegrityError
 from django.test import SimpleTestCase, TestCase
 from django.utils import timezone
+from utils.tests.auth_seam import builtin_auth_service
 
-from account_v2.authentication_plugin_registry import AuthenticationPluginRegistry
 from account_v2.enums import Product
 from account_v2.models import User, UserProductPreference
 from account_v2.product_preference import (
@@ -65,11 +65,9 @@
 
 class _BuiltInSessionAuth(TestCase):
     def setUp(self) -> None:
-        no_plugin = mock.patch.object(
-            AuthenticationPluginRegistry, "is_plugin_available", return_value=False
-        )
-        no_plugin.start()
-        self.addCleanup(no_plugin.stop)
+        # The built-in session auth: a plugin's middleware would 401 a session
+        # it did not seed. Test sessions have no plugin hook yet.
+        self.enterContext(builtin_auth_service())
 
 
 def _user(email: str, **fields) -> User:

backend/account_v2/tests.py

@@ -7,15 +7,14 @@
 
 import secrets
 from types import SimpleNamespace
-from unittest.mock import patch
 
 from django.test import TestCase
 from tenant_account_v2.models import OrganizationMember
+from utils.tests.auth_seam import admin_role, member_role
 from utils.user_context import UserContext
 
 from account_v2.authentication_controller import AuthenticationController
 from account_v2.custom_exceptions import Forbidden
-from account_v2.enums import UserRole
 from account_v2.models import Organization, User
 
 
@@ -32,34 +31,25 @@ def setUp(self) -> None:
             name="org-a", display_name="Org A", organization_id="org-a"
         )
         UserContext.set_organization_identifier(self.org.organization_id)
-        # Pin the built-in auth service: role vocabulary is plugin-defined, so
-        # UserRole below is only meaningful while no auth plugin is installed.
-        registry = patch(
-            "account_v2.authentication_controller."
-            "AuthenticationPluginRegistry.is_plugin_available",
-            return_value=False,
-        )
-        registry.start()
-        self.addCleanup(registry.stop)
         self.controller = AuthenticationController()
-        self.admin = _make_member(self.org, "[email protected]", UserRole.ADMIN.value)
+        self.admin = _make_member(self.org, "[email protected]", admin_role())
 
     def test_sole_admin_cannot_be_demoted(self) -> None:
         with self.assertRaises(Forbidden):
             self.controller._ensure_not_last_admin_demotion(
-                email=self.admin.user.email, new_role=UserRole.USER.value
+                email=self.admin.user.email, new_role=member_role()
             )
 
     def test_demotion_allowed_when_another_admin_remains(self) -> None:
-        _make_member(self.org, "[email protected]", UserRole.ADMIN.value)
+        _make_member(self.org, "[email protected]", admin_role())
         self.controller._ensure_not_last_admin_demotion(
-            email=self.admin.user.email, new_role=UserRole.USER.value
+            email=self.admin.user.email, new_role=member_role()
         )
 
     def test_non_admin_role_change_is_unaffected(self) -> None:
-        member = _make_member(self.org, "[email protected]", UserRole.USER.value)
+        member = _make_member(self.org, "[email protected]", member_role())
         self.controller._ensure_not_last_admin_demotion(
-            email=member.user.email, new_role=UserRole.USER.value
+            email=member.user.email, new_role=member_role()
         )
 
     def test_add_user_role_blocks_sole_admin_demotion(self) -> None:
@@ -69,7 +59,7 @@ def test_add_user_role_blocks_sole_admin_demotion(self) -> None:
                 request,
                 self.org.organization_id,
                 self.admin.user.email,
-                UserRole.USER.value,
+                member_role(),
             )
 
     def test_remove_user_role_blocks_sole_admin_demotion(self) -> None:
@@ -79,5 +69,5 @@ def test_remove_user_role_blocks_sole_admin_demotion(self) -> None:
                 request,
                 self.org.organization_id,
                 self.admin.user.email,
-                UserRole.ADMIN.value,
+                admin_role(),
             )

backend/global_search/tests/test_budget.py

@@ -23,7 +23,6 @@
 from django.db import connection
 from django.test import TestCase
 from django.test.utils import CaptureQueriesContext
-from unittest.mock import patch
 from permissions.roles import ResourceRole
 from pipeline_v2.models import Pipeline
 from prompt_studio.prompt_studio_core_v2.models import CustomTool
@@ -35,6 +34,7 @@
     ResourceMembership,
 )
 from tenant_account_v2.organization_member_service import OrganizationMemberService
+from utils.tests.auth_seam import admin_role, member_role
 from utils.user_context import UserContext
 from workflow_manager.workflow_v2.models.workflow import Workflow
 
@@ -76,8 +76,9 @@ def setUpTestData(cls) -> None:
         users = User.objects.bulk_create(
             User(username=f"u{i}@big.test", email=f"u{i}@big.test") for i in range(USERS)
         )
+        role = member_role()
         OrganizationMember.objects.bulk_create(
-            OrganizationMember(organization=cls.org, user=u, role="user") for u in users
+            OrganizationMember(organization=cls.org, user=u, role=role) for u in users
         )
         groups = OrganizationGroup.objects.bulk_create(
             OrganizationGroup(organization=cls.org, name=f"g{i}", created_by=users[0])
@@ -88,10 +89,10 @@ def setUpTestData(cls) -> None:
         )
         cls.member = users[0]
         cls.admin = users[1]
-        # The role lookup a request pays is counted: only the role-name check
-        # is patched (see setUp), never the predicate that queries it.
+        # The role lookup a request pays is counted: nothing on the admin
+        # predicate is patched.
         OrganizationMember.objects.filter(organization=cls.org, user=cls.admin).update(
-            role="admin"
+            role=admin_role()
         )
 
         common = {"organization": cls.org}
@@ -101,7 +102,9 @@ def setUpTestData(cls) -> None:
         )
         rows = [
             *Workflow.objects.bulk_create(
-                Workflow(workflow_name=_name("wf", i), description=_name("d", i), **common)
+                Workflow(
+                    workflow_name=_name("wf", i), description=_name("d", i), **common
+                )
                 for i in n
             ),
             *CustomTool.objects.bulk_create(
@@ -146,7 +149,9 @@ def setUpTestData(cls) -> None:
                 for i in n
             ),
             *ConnectorInstance.objects.bulk_create(
-                ConnectorInstance(connector_name=_name("conn", i), connector_id="minio|seed", **common)
+                ConnectorInstance(
+                    connector_name=_name("conn", i), connector_id="minio|seed", **common
+                )
                 for i in n
             ),
         ]
@@ -177,7 +182,14 @@ def setUpTestData(cls) -> None:
                 row.shared_to_org = True
         ResourceMembership.objects.bulk_create(memberships, batch_size=2_000)
         ResourceGroupShare.objects.bulk_create(shares, batch_size=2_000)
-        for model in (Workflow, CustomTool, APIDeployment, Pipeline, AdapterInstance, ConnectorInstance):
+        for model in (
+            Workflow,
+            CustomTool,
+            APIDeployment,
+            Pipeline,
+            AdapterInstance,
+            ConnectorInstance,
+        ):
             model.objects.filter(
                 pk__in=[r.pk for r in rows if isinstance(r, model) and r.shared_to_org]
             ).update(shared_to_org=True)
@@ -186,17 +198,6 @@ def setUpTestData(cls) -> None:
 
     def setUp(self) -> None:
         UserContext.set_organization_identifier(self.org.organization_id)
-        # Each auth plugin names the admin role differently ("admin" in OSS,
-        # "unstract_admin" in cloud). Pin the name check, which is pure
-        # Python, so the admin subtest measures an admin in every
-        # distribution; the membership query still runs and is counted.
-        name_check = patch(
-            "account_v2.authentication_controller.AuthenticationController"
-            ".is_admin_by_role",
-            new=lambda _self, role: role == "admin",
-        )
-        name_check.start()
-        self.addCleanup(name_check.stop)
         self.assertTrue(
             OrganizationMemberService.is_user_organization_admin(
                 User.objects.get(pk=self.admin.pk)
@@ -205,7 +206,8 @@ def setUp(self) -> None:
 
     def _queries(self, user: User, limit: int) -> list[str]:
         """SQL issued for one search, minus the per-provider savepoints the
-        view opens inside a transaction (TestCase always has one)."""
+        view opens inside a transaction (TestCase always has one).
+        """
         # A fresh user object, as each request gets: the admin-role memo
         # lives on the user, so reusing one would hide that lookup.
         user = User.objects.get(pk=user.pk)
@@ -254,6 +256,8 @@ def test_p95_is_measured(self) -> None:
         for user in (self.member, self.admin):
             with self.subTest(user=user.email):
                 p95 = self._p95_ms(user)
-                print(f"[global_search budget] {user.email} p95={p95:.1f}ms (budget {BUDGET_MS}ms)")
+                print(
+                    f"[global_search budget] {user.email} p95={p95:.1f}ms (budget {BUDGET_MS}ms)"
+                )
                 if ASSERT_P95:
                     self.assertLessEqual(p95, CEILING_MS)

backend/global_search/tests/test_contract.py

@@ -14,16 +14,17 @@
 from unittest.mock import patch
 
 import jsonschema
-from django.db import connection
 from account_v2.models import Organization
+from django.db import connection
 from django.test import TestCase
 from django.urls import resolve
 from permissions.roles import ResourceRole
 from permissions.tests.base import CoOwnerOrgTestMixin, make_user
 from tenant_account_v2.models import OrganizationMember
+from utils.tests.auth_seam import admin_role, builtin_auth_service
 
 from global_search.constants import OSSSearchType as T
-from global_search.oss_providers import AdapterProvider, OSS_PROVIDERS, WorkflowProvider
+from global_search.oss_providers import OSS_PROVIDERS, AdapterProvider, WorkflowProvider
 from global_search.providers import SearchProvider
 from global_search.tests.base import OSS_BUILDERS, search
 from global_search.views import GlobalSearchView
@@ -136,16 +137,15 @@ def test_url_org_must_match_the_session_org(self) -> None:
         other = Organization.objects.create(
             name="org-b", display_name="Org B", organization_id="org-b"
         )
-        OrganizationMember.objects.create(organization=other, user=self.admin, role="admin")
         # The built-in session auth, so a distribution's auth plugin cannot
-        # 401 the request first; the org check runs after either.
-        no_plugin = patch(
-            "account_v2.custom_auth_middleware.AuthenticationPluginRegistry"
-            ".is_plugin_available",
-            return_value=False,
+        # 401 the request first; the org check runs after either. Test
+        # sessions have no plugin hook yet. The admin's rows then take the
+        # built-in service's role name.
+        self.enterContext(builtin_auth_service())
+        OrganizationMember.objects.filter(user=self.admin).update(role=admin_role())
+        OrganizationMember.objects.create(
+            organization=other, user=self.admin, role=admin_role()
         )
-        no_plugin.start()
-        self.addCleanup(no_plugin.stop)
         self.client.force_login(self.admin)
         session = self.client.session
         session["organization"] = self.org.organization_id

backend/permissions/tests/base.py

@@ -9,23 +9,16 @@
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
-from unittest.mock import patch
 
 from account_v2.models import Organization, User
-from permissions.permission import IsOwner
 from rest_framework.test import APIRequestFactory
 from rest_framework.views import APIView
 from tenant_account_v2.models import OrganizationMember
+from utils.tests.auth_seam import admin_role, member_role
 from utils.user_context import UserContext
 from workflow_manager.workflow_v2.models.workflow import Workflow
 
-# Admin resolution is auth-plugin backed, so patch it to a deterministic
-# predicate — the ``role="admin"`` membership row alone does not make
-# ``_is_organization_admin`` true.
-_ADMIN_PREDICATE = (
-    "tenant_account_v2.organization_member_service."
-    "OrganizationMemberService.is_user_organization_admin"
-)
+from permissions.permission import IsOwner
 
 
 def make_user(email: str, **kwargs: Any) -> User:
@@ -120,7 +113,7 @@ class ResourceSpec:
 
 
 class CoOwnerOrgTestMixin:
-    """setUp helper: one org, members, and a deterministic admin predicate."""
+    """setUp helper: one org and its members, admin by the active auth service's role."""
 
     def _seed_org(self) -> None:
         self.org = Organization.objects.create(
@@ -135,18 +128,11 @@ def _seed_org(self) -> None:
         self.stranger = make_user("[email protected]")  # NOT an org member
         for user in (self.owner, self.coowner, self.viewer, self.outsider):
             OrganizationMember.objects.create(
-                organization=self.org, user=user, role="user"
+                organization=self.org, user=user, role=member_role()
             )
         OrganizationMember.objects.create(
-            organization=self.org, user=self.admin, role="admin"
-        )
-        patcher = patch(
-            _ADMIN_PREDICATE,
-            side_effect=lambda u: getattr(u, "email", None) == self.admin.email,
+            organization=self.org, user=self.admin, role=admin_role()
         )
-        patcher.start()
-        # addCleanup is provided by TestCase (this is a mixin combined with it).
-        self.addCleanup(patcher.stop)  # type: ignore[attr-defined]
 
     def _is_owner_perm(self, user: User, obj: object) -> bool:
         """Run ``IsOwner.has_object_permission`` for ``user`` against ``obj``."""