Add tests for new content type RBAC model

Author: AlanCoding

ansible_base/rbac/management/__init__.py

@@ -12,7 +12,7 @@
 def create_dab_permissions(app_config, verbosity=2, interactive=True, using=DEFAULT_DB_ALIAS, apps=global_apps, **kwargs):
     """
     This is modified from the django auth.
-    This will create DABPermission entries
+    This will create DABPermission and DABContentType entries
     this will only create permissions for registered models
     """
     if not getattr(app_config, 'models_module', None):
@@ -23,6 +23,7 @@ def create_dab_permissions(app_config, verbosity=2, interactive=True, using=DEFA
     if not any(cls._meta.app_label == app_label for cls in permission_registry._registry):
         return
 
+    # TODO: remove when migration is finished
     # Ensure that contenttypes are created for this app. Needed if
     # 'ansible_base.rbac' is in INSTALLED_APPS before
     # 'django.contrib.contenttypes'.
@@ -45,13 +46,21 @@ def create_dab_permissions(app_config, verbosity=2, interactive=True, using=DEFA
     if not router.allow_migrate_model(using, Permission):
         return
 
+    rbac_models = [klass for klass in app_config.get_models() if permission_registry.is_registered(klass)]
+
+    if not rbac_models:
+        logger.debug(f'No RBAC models registered for app {app_label}')
+        return
+
+    from .create_types import create_DAB_contenttypes
+
+    create_DAB_contenttypes(rbac_models, verbosity=verbosity, using=using, apps=apps)
+
     # This will hold the permissions we're looking for as (content_type, (codename, name))
     searched_perms = []
     # The codenames and ctypes that should exist.
     ctypes = set()
-    for klass in app_config.get_models():
-        if not permission_registry.is_registered(klass):
-            continue
+    for klass in rbac_models:
         # Force looking up the content types in the current database
         # before creating foreign keys to them.
         ctype = ContentType.objects.db_manager(using).get_for_model(klass, for_concrete_model=False)

ansible_base/rbac/management/create_types.py

@@ -0,0 +1,59 @@
+from typing import Type
+
+from django.apps import apps as global_apps
+from django.db import DEFAULT_DB_ALIAS, models
+
+
+def get_local_DAB_contenttypes(using: str, ct_model: Type[models.Model], service: str) -> dict[tuple[str, str], models.Model]:
+    # This should work in migration scenarios, but other code checks for existence of it on manager
+    ct_model.objects.clear_cache()
+
+    return {
+        (service, ct.model): ct
+        for ct in ct_model.objects.using(using).filter(service=service)
+    }
+
+
+def create_DAB_contenttypes(
+    rbac_models: list[Type[models.Model]],
+    verbosity=2,
+    using=DEFAULT_DB_ALIAS,
+    apps=global_apps
+):
+    """Create DABContentType for models in the given app.
+
+    This is significantly different from the ContentType post-migrate method,
+    because that creates types for all apps, and so this is only called one app at a time.
+    DAB RBAC runs its post_migration logic just once, because the model list
+    comes from the permission registry.
+    """
+    DABContentType = apps.get_model("dab_rbac", "DABContentType")
+
+    from ansible_base.rbac.remote import get_local_resource_prefix
+
+    service = get_local_resource_prefix()
+
+    content_types = get_local_DAB_contenttypes(
+        using, DABContentType, service
+    )
+    if not rbac_models:
+        return
+
+    cts = [
+        DABContentType(
+            service=service,
+            app_label=model._meta.app_label,
+            model=model._meta.model_name,
+        )
+        for model in rbac_models
+        if (service, model._meta.model_name) not in content_types
+    ]
+    if not cts:
+        return
+    DABContentType.objects.using(using).bulk_create(cts)
+    if verbosity >= 2:
+        for ct in cts:
+            print(
+                "Adding DAB content type "
+                f"'{ct.service}:{ct.app_label} | {ct.model}'"
+            )

ansible_base/rbac/models/content_type.py

@@ -6,7 +6,7 @@
 from django.db.models.options import Options
 from django.utils.translation import gettext_lazy as _
 
-from ..remote import get_remote_object_class, get_local_resource_prefix
+from ..remote import get_remote_object_class, get_local_resource_prefix, RemoteObject
 
 
 class DABContentTypeManager(django_models.Manager["DABContentType"]):
@@ -191,7 +191,7 @@ def model_class(self) -> Optional[Type[django_models.Model]]:
         except LookupError:
             return None
 
-    def get_object_for_this_type(self, **kwargs: Any) -> django_models.Model:
+    def get_object_for_this_type(self, **kwargs: Any) -> django_models.Model | RemoteObject:
         """Return the object referenced by this content type."""
         model = self.model_class()
         if model is None:
@@ -201,7 +201,7 @@ def get_object_for_this_type(self, **kwargs: Any) -> django_models.Model:
             return get_remote_object_class()(self, object_id)
         return model._base_manager.get(**kwargs)
 
-    def get_all_objects_for_this_type(self, **kwargs: Any) -> django_models.QuerySet | Sequence[django_models.Model]:
+    def get_all_objects_for_this_type(self, **kwargs: Any) -> django_models.QuerySet | Sequence[django_models.Model | RemoteObject]:
         """Return all objects referenced by this content type."""
         model = self.model_class()
         if model is None:

ansible_base/rbac/remote.py

@@ -39,7 +39,7 @@ def get_local_resource_prefix() -> str:
     return 'local'
 
 
-def get_resource_prefix(cls: models.Model) -> str:
+def get_resource_prefix(cls: Type[models.Model]) -> str:
     """The API project designator for given cls, according to the resource registry
 
     This is used for related slug references, like "awx.inventory" to reference

ansible_base/rbac/validators.py

@@ -9,6 +9,8 @@
 from ansible_base.lib.utils.models import is_add_perm
 from ansible_base.rbac.permission_registry import permission_registry
 
+from .remote import get_resource_prefix
+
 
 def system_roles_enabled():
     return bool(
@@ -163,7 +165,7 @@ def validate_permissions_for_model(permissions, content_type: Optional[Model], m
             if content_type and perm.codename.startswith('view'):
                 continue
             model = perm.content_type.model_class()
-            if permission_registry.get_resource_prefix(model) == 'shared':
+            if get_resource_prefix(model) == 'shared':
                 raise ValidationError({'permissions', 'Local custom roles can only include view permission for shared models'})
 
 

test_app/tests/rbac/models/test_content_type.py

@@ -0,0 +1,82 @@
+import pytest
+from django.db import models
+from django.test import TestCase
+from django.test.utils import isolate_apps
+
+from ansible_base.rbac.models import DABContentType
+from ansible_base.rbac.remote import RemoteObject
+
+from test_app.models import Inventory
+
+
+@pytest.mark.django_db
+def test_post_migrate_creates_contenttype():
+    ct = DABContentType.objects.get(app_label="test_app", model="inventory")
+    assert ct.service == "aap"
+
+
+@pytest.mark.django_db
+class DABContentTypeTests(TestCase):
+    """These tests originally came from Django contenttypes"""
+    def setUp(self):
+        DABContentType.objects.clear_cache()
+        self.addCleanup(DABContentType.objects.clear_cache)
+
+    def test_lookup_cache(self):
+        with self.assertNumQueries(1):
+            DABContentType.objects.get_for_model(Inventory)
+        with self.assertNumQueries(0):
+            ct = DABContentType.objects.get_for_model(Inventory)
+        with self.assertNumQueries(0):
+            DABContentType.objects.get_for_id(ct.id)
+        with self.assertNumQueries(0):
+            DABContentType.objects.get_by_natural_key(
+                ct.service,
+                ct.app_label,
+                ct.model,
+            )
+        DABContentType.objects.clear_cache()
+        with self.assertNumQueries(1):
+            DABContentType.objects.get_for_model(Inventory)
+
+    @isolate_apps("tests")
+    def test_get_for_model_create_contenttype(self):
+        class ModelCreatedOnTheFly(models.Model):
+            name = models.CharField(max_length=10)
+
+            class Meta:
+                app_label = "tests"
+
+        ct = DABContentType.objects.get_for_model(ModelCreatedOnTheFly)
+        assert ct.app_label == "tests"
+        assert ct.model == "modelcreatedonthefly"
+
+
+@pytest.mark.django_db
+def test_get_object_for_this_type_remote():
+    """Remote objects should return a remote proxy."""
+    ct = DABContentType.objects.create(
+        service="remote_proj",
+        app_label="testapp",
+        model="book",
+    )
+
+    obj = ct.get_object_for_this_type(pk=1)
+
+    assert isinstance(obj, RemoteObject)
+    assert obj.object_id == 1
+    assert obj.content_type == ct
+
+
+@pytest.mark.django_db
+def test_get_all_objects_for_this_type_remote():
+    ct = DABContentType.objects.create(
+        service="remote_proj2",
+        app_label="testapp",
+        model="book",
+    )
+
+    objs = ct.get_all_objects_for_this_type(pk__in=[1, 2])
+
+    assert [o.object_id for o in objs] == [1, 2]
+    assert all(isinstance(o, RemoteObject) for o in objs)