refactor(rbac): drop test singular scope seed

jordan-umusu · jordan-umusu · commit ed1dadb193a6 · 2026-02-09T16:03:51.000-05:00
diff --git a/tests/unit/test_authz_seeding.py b/tests/unit/test_authz_seeding.py
@@ -1,16 +1,19 @@
 """Tests for RBAC scope seeding."""
 
+from uuid import uuid4
+
 import pytest
-from sqlalchemy import select
+from sqlalchemy import func, select
 
 from tracecat.authz.enums import ScopeSource
 from tracecat.authz.seeding import (
+    PRESET_ROLE_DEFINITIONS,
     SYSTEM_SCOPE_DEFINITIONS,
-    seed_registry_scope,
-    seed_registry_scopes_bulk,
+    seed_registry_scopes,
+    seed_system_roles_for_all_orgs,
     seed_system_scopes,
 )
-from tracecat.db.models import Scope
+from tracecat.db.models import Organization, Role, RoleScope, Scope
 
 
 @pytest.mark.anyio
@@ -61,41 +64,7 @@ async def test_seed_system_scopes_idempotent(session):
 
 
 @pytest.mark.anyio
-async def test_seed_registry_scope(session):
-    """Test seeding a single registry scope."""
-    action_key = "tools.test_integration.test_action"
-
-    scope = await seed_registry_scope(session, action_key, "Test action scope")
-    await session.commit()
-
-    assert scope is not None
-    assert scope.name == f"action:{action_key}:execute"
-    assert scope.resource == "action"
-    assert scope.action == "execute"
-    assert scope.source == ScopeSource.PLATFORM
-    assert scope.source_ref == action_key
-    assert scope.organization_id is None
-
-
-@pytest.mark.anyio
-async def test_seed_registry_scope_idempotent(session):
-    """Test that seeding the same registry scope twice returns existing scope."""
-    action_key = "tools.test_integration.test_action_idempotent"
-
-    # Seed first time
-    scope1 = await seed_registry_scope(session, action_key)
-    await session.commit()
-
-    # Seed second time
-    scope2 = await seed_registry_scope(session, action_key)
-
-    assert scope1 is not None
-    assert scope2 is not None
-    assert scope1.id == scope2.id
-
-
-@pytest.mark.anyio
-async def test_seed_registry_scopes_bulk(session):
+async def test_seed_registry_scopes(session):
     """Test bulk seeding of registry scopes."""
     action_keys = [
         "tools.okta.list_users",
@@ -104,7 +73,7 @@ async def test_seed_registry_scopes_bulk(session):
         "core.http_request",
     ]
 
-    inserted_count = await seed_registry_scopes_bulk(session, action_keys)
+    inserted_count = await seed_registry_scopes(session, action_keys)
     await session.commit()
 
     assert inserted_count == len(action_keys)
@@ -125,28 +94,91 @@ async def test_seed_registry_scopes_bulk(session):
 
 
 @pytest.mark.anyio
-async def test_seed_registry_scopes_bulk_idempotent(session):
+async def test_seed_registry_scopes_idempotent(session):
     """Test that bulk seeding is idempotent."""
     action_keys = ["tools.test.action1", "tools.test.action2"]
 
     # First seed
-    first_count = await seed_registry_scopes_bulk(session, action_keys)
+    first_count = await seed_registry_scopes(session, action_keys)
     await session.commit()
     assert first_count == 2
 
     # Second seed
-    second_count = await seed_registry_scopes_bulk(session, action_keys)
+    second_count = await seed_registry_scopes(session, action_keys)
     await session.commit()
     assert second_count == 0
 
 
 @pytest.mark.anyio
-async def test_seed_registry_scopes_bulk_empty(session):
+async def test_seed_registry_scopes_empty(session):
     """Test bulk seeding with empty list."""
-    inserted_count = await seed_registry_scopes_bulk(session, [])
+    inserted_count = await seed_registry_scopes(session, [])
     assert inserted_count == 0
 
 
+@pytest.mark.anyio
+async def test_seed_system_roles_for_all_orgs_creates_roles_and_links(session):
+    """Seed preset roles for all orgs and link expected system scopes."""
+    # Ensure scope IDs exist for role->scope links.
+    await seed_system_scopes(session)
+
+    # Add an extra org so the function processes multiple orgs in one call.
+    extra_org = Organization(
+        id=uuid4(),
+        name="Extra test org",
+        slug=f"extra-test-org-{uuid4().hex[:8]}",
+        is_active=True,
+    )
+    session.add(extra_org)
+    await session.flush()
+
+    # Capture target org IDs in this isolated session.
+    org_result = await session.execute(select(Organization.id))
+    org_ids = {org_id for (org_id,) in org_result.tuples().all()}
+
+    created_by_org = await seed_system_roles_for_all_orgs(session)
+
+    for org_id in org_ids:
+        assert created_by_org[org_id] == len(PRESET_ROLE_DEFINITIONS)
+
+    roles_result = await session.execute(
+        select(Role.id, Role.organization_id, Role.slug).where(
+            Role.organization_id.in_(org_ids),
+            Role.slug.in_(PRESET_ROLE_DEFINITIONS),
+        )
+    )
+    roles = roles_result.tuples().all()
+    assert len(roles) == len(org_ids) * len(PRESET_ROLE_DEFINITIONS)
+
+    role_scope_count_stmt = (
+        select(Role.slug, func.count(RoleScope.scope_id))
+        .select_from(Role)
+        .join(RoleScope, RoleScope.role_id == Role.id)
+        .where(
+            Role.organization_id.in_(org_ids),
+            Role.slug.in_(PRESET_ROLE_DEFINITIONS),
+        )
+        .group_by(Role.slug)
+    )
+    role_scope_count_result = await session.execute(role_scope_count_stmt)
+    role_scope_counts = dict(role_scope_count_result.tuples().all())
+    expected_org_count = len(org_ids)
+    for slug, role_def in PRESET_ROLE_DEFINITIONS.items():
+        assert role_scope_counts[slug] == len(role_def.scopes) * expected_org_count
+
+
+@pytest.mark.anyio
+async def test_seed_system_roles_for_all_orgs_idempotent(session):
+    """Running role seeding twice should not create duplicate roles."""
+    await seed_system_scopes(session)
+
+    first = await seed_system_roles_for_all_orgs(session)
+    second = await seed_system_roles_for_all_orgs(session)
+
+    assert sum(first.values()) > 0
+    assert all(created == 0 for created in second.values())
+
+
 @pytest.mark.anyio
 async def test_system_scope_definitions_format(session):
     """Test that all system scope definitions follow the expected format."""
diff --git a/tracecat/authz/seeding.py b/tracecat/authz/seeding.py
@@ -243,7 +243,7 @@ async def seed_system_scopes(session: AsyncSession) -> int:
     result = await session.execute(stmt)
     await session.commit()
 
-    inserted_count = result.rowcount if result.rowcount else 0  # pyright: ignore[reportAttributeAccessIssue]
+    inserted_count = result.rowcount or 0  # pyright: ignore[reportAttributeAccessIssue]
     logger.info(
         "System scopes seeded",
         inserted=inserted_count,
@@ -252,61 +252,7 @@ async def seed_system_scopes(session: AsyncSession) -> int:
     return inserted_count
 
 
-async def seed_registry_scope(
-    session: AsyncSession,
-    action_key: str,
-    description: str | None = None,
-) -> Scope | None:
-    """Seed a single registry action scope.
-
-    Creates a scope for `action:{action_key}:execute` if it doesn't exist.
-    Registry scopes have organization_id=NULL and source='registry'.
-    Uses upsert (ON CONFLICT DO NOTHING) for concurrency safety.
-
-    Args:
-        session: Database session
-        action_key: The action key (e.g., "tools.okta.list_users")
-        description: Optional description for the scope
-
-    Returns:
-        The created or existing Scope
-    """
-    scope_name = f"action:{action_key}:execute"
-    scope_id = uuid4()
-
-    # Use upsert for concurrency safety
-    stmt = pg_insert(Scope).values(
-        id=scope_id,
-        name=scope_name,
-        resource="action",
-        action="execute",
-        description=description or f"Execute {action_key} action",
-        source=ScopeSource.PLATFORM,
-        source_ref=action_key,
-        organization_id=None,
-    )
-    stmt = stmt.on_conflict_do_nothing(
-        index_elements=["name"], index_where=Scope.organization_id.is_(None)
-    )
-    result = await session.execute(stmt)
-    await session.flush()
-
-    # Re-query to get the scope (whether newly inserted or already existing)
-    select_stmt = select(Scope).where(
-        Scope.name == scope_name, Scope.organization_id.is_(None)
-    )
-    select_result = await session.execute(select_stmt)
-    scope = select_result.scalar_one_or_none()
-
-    if result.rowcount and result.rowcount > 0:  # pyright: ignore[reportAttributeAccessIssue]
-        logger.debug(
-            "Registry scope created", scope_name=scope_name, action_key=action_key
-        )
-
-    return scope
-
-
-async def seed_registry_scopes_bulk(
+async def seed_registry_scopes(
     session: AsyncSession,
     action_keys: list[str],
 ) -> int:
@@ -413,7 +359,7 @@ async def seed_system_roles_for_org(
         index_elements=["organization_id", "slug"]
     )
     result = await session.execute(role_stmt)
-    roles_created = result.rowcount if result.rowcount else 0  # pyright: ignore[reportAttributeAccessIssue]
+    roles_created = result.rowcount or 0  # pyright: ignore[reportAttributeAccessIssue]
 
     # Re-query to get actual role IDs (may differ if roles already existed)
     existing_roles_stmt = select(Role.id, Role.slug).where(
@@ -490,14 +436,73 @@ async def seed_system_roles_for_all_orgs(session: AsyncSession) -> dict[UUID, in
         logger.info("No organizations found, skipping system role seeding")
         return {}
 
-    results: dict[UUID, int] = {}
-    total_created = 0
-
+    # 1) Upsert all preset roles for all orgs in one bulk query.
+    role_values = []
     for org_id in org_ids:
-        roles_created = await seed_system_roles_for_org(session, org_id)
-        results[org_id] = roles_created
-        total_created += roles_created
+        for slug, role_def in PRESET_ROLE_DEFINITIONS.items():
+            role_values.append(
+                {
+                    "id": uuid4(),
+                    "name": role_def.name,
+                    "slug": slug,
+                    "description": role_def.description,
+                    "organization_id": org_id,
+                    "created_by": None,
+                }
+            )
+
+    role_insert_stmt = pg_insert(Role).values(role_values)
+    role_insert_stmt = role_insert_stmt.on_conflict_do_nothing(
+        index_elements=["organization_id", "slug"]
+    ).returning(Role.organization_id)
+    role_insert_result = await session.execute(role_insert_stmt)
+
+    # Return shape remains {org_id: roles_created_for_org}.
+    results: dict[UUID, int] = dict.fromkeys(org_ids, 0)
+    for (organization_id,) in role_insert_result.tuples().all():
+        results[organization_id] += 1
+
+    # 2) Fetch all relevant global scopes once.
+    scope_stmt = select(Scope.id, Scope.name).where(Scope.organization_id.is_(None))
+    scope_result = await session.execute(scope_stmt)
+    scope_id_by_name: dict[str, UUID] = {
+        scope_name: scope_id for scope_id, scope_name in scope_result.tuples().all()
+    }
+
+    # 3) Fetch all preset roles for those orgs and bulk insert role-scope links.
+    role_stmt = select(Role.id, Role.slug).where(
+        Role.organization_id.in_(org_ids),
+        Role.slug.in_(PRESET_ROLE_DEFINITIONS),
+    )
+    role_result = await session.execute(role_stmt)
+
+    role_scope_values = []
+    for role_id, role_slug in role_result.tuples().all():
+        if role_slug is None:
+            continue
+        role_def = PRESET_ROLE_DEFINITIONS.get(role_slug)
+        if role_def is None:
+            continue
+        for scope_name in role_def.scopes:
+            scope_id = scope_id_by_name.get(scope_name)
+            if scope_id is None:
+                logger.warning(
+                    "Scope not found for system role",
+                    scope_name=scope_name,
+                    role_slug=role_slug,
+                )
+                continue
+            role_scope_values.append({"role_id": role_id, "scope_id": scope_id})
+
+    if role_scope_values:
+        role_scope_stmt = pg_insert(RoleScope).values(role_scope_values)
+        role_scope_stmt = role_scope_stmt.on_conflict_do_nothing(
+            index_elements=["role_id", "scope_id"]
+        )
+        await session.execute(role_scope_stmt)
+    await session.commit()
 
+    total_created = sum(results.values())
     logger.info(
         "System roles seeded for all organizations",
         num_orgs=len(org_ids),
diff --git a/tracecat/registry/sync/jobs.py b/tracecat/registry/sync/jobs.py
@@ -12,7 +12,7 @@
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from tracecat.authz.seeding import seed_registry_scopes_bulk
+from tracecat.authz.seeding import seed_registry_scopes
 from tracecat.db.engine import get_async_session_context_manager
 from tracecat.db.locks import (
     derive_lock_key_from_parts,
@@ -195,7 +195,7 @@ async def _seed_registry_scopes(
     action_keys = [f"{action.namespace}.{action.name}" for action in actions]
 
     try:
-        inserted = await seed_registry_scopes_bulk(session, action_keys)
+        inserted = await seed_registry_scopes(session, action_keys)
         await session.commit()
         logger.info("Registry scopes seeded", inserted=inserted, total=len(action_keys))
     except DBAPIError as e:
