diff --git a/openedx_learning/apps/authoring/publishing/api.py b/openedx_learning/apps/authoring/publishing/api.py
index 895f091fd..5fece0ac8 100644
--- a/openedx_learning/apps/authoring/publishing/api.py
+++ b/openedx_learning/apps/authoring/publishing/api.py
@@ -13,12 +13,7 @@
 from django.db.models import F, Q, QuerySet
 from django.db.transaction import atomic
 
-from .model_mixins import (
-    ContainerMixin,
-    PublishableContentModelRegistry,
-    PublishableEntityMixin,
-    PublishableEntityVersionMixin,
-)
+from .model_mixins import PublishableContentModelRegistry, PublishableEntityMixin, PublishableEntityVersionMixin
 from .models import (
     Container,
     ContainerVersion,
@@ -579,6 +574,7 @@ def get_published_version_as_of(entity_id: int, publish_log_id: int) -> Publisha
 def create_container(
     learning_package_id: int,
     key: str,
+    container_type: str,
     created: datetime,
     created_by: int | None,
 ) -> Container:
@@ -595,12 +591,14 @@ def create_container(
     Returns:
         The newly created container.
     """
+    assert container_type  # Shouldn't be empty/none
     with atomic():
         publishable_entity = create_publishable_entity(
             learning_package_id, key, created, created_by
         )
         container = Container.objects.create(
             publishable_entity=publishable_entity,
+            container_type=container_type,
         )
     return container
 
@@ -635,7 +633,7 @@ def create_entity_list_with_rows(
         The newly created entity list.
     """
     order_nums = range(len(entity_pks))
-    with atomic():
+    with atomic(savepoint=False):
         entity_list = create_entity_list()
         EntityListRow.objects.bulk_create(
             [
@@ -679,7 +677,7 @@ def create_container_version(
     Returns:
         The newly created container version.
     """
-    with atomic():
+    with atomic(savepoint=False):
         container = Container.objects.select_related("publishable_entity").get(pk=container_pk)
         entity = container.publishable_entity
 
@@ -789,6 +787,7 @@ def create_container_and_version(
     learning_package_id: int,
     key: str,
     *,
+    container_type: str,
     created: datetime,
     created_by: int | None,
     title: str,
@@ -812,7 +811,7 @@ def create_container_and_version(
         The newly created container version.
     """
     with atomic():
-        container = create_container(learning_package_id, key, created, created_by)
+        container = create_container(learning_package_id, key, container_type, created, created_by)
         container_version = create_container_version(
             container.publishable_entity.pk,
             1,
@@ -854,15 +853,13 @@ def entity(self):
 
 
 def get_entities_in_draft_container(
-    container: Container | ContainerMixin,
+    container: Container,
 ) -> list[ContainerEntityListEntry]:
     """
     [ 🛑 UNSTABLE ]
     Get the list of entities and their versions in the draft version of the
     given container.
     """
-    if isinstance(container, ContainerMixin):
-        container = container.container
     assert isinstance(container, Container)
     entity_list = []
     for row in container.versioning.draft.entity_list.entitylistrow_set.order_by("order_num"):
@@ -877,19 +874,15 @@ def get_entities_in_draft_container(
 
 
 def get_entities_in_published_container(
-    container: Container | ContainerMixin,
+    container: Container,
 ) -> list[ContainerEntityListEntry] | None:
     """
     [ 🛑 UNSTABLE ]
     Get the list of entities and their versions in the published version of the
     given container.
     """
-    if isinstance(container, ContainerMixin):
-        cv = container.container.versioning.published
-    elif isinstance(container, Container):
-        cv = container.versioning.published
-    else:
-        raise TypeError(f"Expected Container or ContainerMixin; got {type(container)}")
+    assert isinstance(container, Container)
+    cv = container.versioning.published
     if cv is None:
         return None  # There is no published version of this container. Should this be an exception?
     assert isinstance(cv, ContainerVersion)
@@ -905,9 +898,7 @@ def get_entities_in_published_container(
     return entity_list
 
 
-def contains_unpublished_changes(
-    container: Container | ContainerMixin,
-) -> bool:
+def contains_unpublished_changes(container_id: int) -> bool:
     """
     [ 🛑 UNSTABLE ]
     Check recursively if a container has any unpublished changes.
@@ -920,14 +911,10 @@ def contains_unpublished_changes(
     that's in the container, it will be `False`. This method will return `True`
     in either case.
     """
-    if isinstance(container, ContainerMixin):
-        # This is similar to 'get_container(container.container_id)' but pre-loads more data.
-        container = Container.objects.select_related(
-            "publishable_entity__draft__version__containerversion__entity_list",
-        ).get(pk=container.container_id)
-    else:
-        pass  # TODO: select_related if we're given a raw Container rather than a ContainerMixin like Unit?
-    assert isinstance(container, Container)
+    # This is similar to 'get_container(container.container_id)' but pre-loads more data.
+    container = Container.objects.select_related(
+        "publishable_entity__draft__version__containerversion__entity_list",
+    ).get(pk=container_id)
 
     if container.versioning.has_unpublished_changes:
         return True
@@ -949,7 +936,7 @@ def contains_unpublished_changes(
             child_container = None
         if child_container:
             # This is itself a container - check recursively:
-            if contains_unpublished_changes(child_container):
+            if contains_unpublished_changes(child_container.pk):
                 return True
         else:
             # This is not a container:
diff --git a/openedx_learning/apps/authoring/publishing/migrations/0003_containers.py b/openedx_learning/apps/authoring/publishing/migrations/0003_containers.py
index 7260e2a0f..8048017e5 100644
--- a/openedx_learning/apps/authoring/publishing/migrations/0003_containers.py
+++ b/openedx_learning/apps/authoring/publishing/migrations/0003_containers.py
@@ -1,8 +1,10 @@
-# Generated by Django 4.2.19 on 2025-03-07 23:09
+# Generated by Django 4.2.19 on 2025-03-10 18:03
 
 import django.db.models.deletion
 from django.db import migrations, models
 
+import openedx_learning.lib.fields
+
 
 class Migration(migrations.Migration):
 
@@ -15,6 +17,7 @@ class Migration(migrations.Migration):
             name='Container',
             fields=[
                 ('publishable_entity', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, primary_key=True, serialize=False, to='oel_publishing.publishableentity')),
+                ('container_type', openedx_learning.lib.fields.MultiCollationCharField(db_collations={'mysql': 'utf8mb4_bin', 'sqlite': 'BINARY'}, max_length=500)),
             ],
             options={
                 'abstract': False,
diff --git a/openedx_learning/apps/authoring/publishing/model_mixins/__init__.py b/openedx_learning/apps/authoring/publishing/model_mixins/__init__.py
index 47f902970..14ea9315a 100644
--- a/openedx_learning/apps/authoring/publishing/model_mixins/__init__.py
+++ b/openedx_learning/apps/authoring/publishing/model_mixins/__init__.py
@@ -1,5 +1,4 @@
 """
 Mixins provided by the publishing app
 """
-from .container import *
 from .publishable_entity import *
diff --git a/openedx_learning/apps/authoring/publishing/model_mixins/container.py b/openedx_learning/apps/authoring/publishing/model_mixins/container.py
deleted file mode 100644
index 944d49ac0..000000000
--- a/openedx_learning/apps/authoring/publishing/model_mixins/container.py
+++ /dev/null
@@ -1,98 +0,0 @@
-"""
-ContainerMixin and ContainerVersionMixin
-"""
-from __future__ import annotations
-
-from datetime import datetime
-from typing import TYPE_CHECKING, ClassVar, Self
-
-from django.db import models
-
-from openedx_learning.lib.managers import WithRelationsManager
-
-from .publishable_entity import PublishableEntityMixin, PublishableEntityVersionMixin
-
-if TYPE_CHECKING:
-    from ..models.container import Container, ContainerVersion
-else:
-    # To avoid circular imports, we need to reference these models using strings only
-    Container = "oel_publishing.Container"
-    ContainerVersion = "oel_publishing.ContainerVersion"
-
-__all__ = [
-    "ContainerMixin",
-    "ContainerVersionMixin",
-]
-
-
-class ContainerMixin(PublishableEntityMixin):
-    """
-    Convenience mixin to link your models against Container.
-
-    Please see docstring for Container for more details.
-
-    If you use this class, you *MUST* also use ContainerVersionMixin
-    """
-
-    # select these related entities by default for all queries
-    objects: ClassVar[WithRelationsManager[Self]] = WithRelationsManager(  # type: ignore[assignment]
-        "container",
-        "publishable_entity",
-        "publishable_entity__published",
-        "publishable_entity__draft",
-    )
-
-    container = models.OneToOneField(
-        Container,
-        on_delete=models.CASCADE,
-    )
-
-    @property
-    def uuid(self) -> str:
-        return self.container.uuid
-
-    @property
-    def created(self) -> datetime:
-        return self.container.created
-
-    class Meta:
-        abstract = True
-
-
-class ContainerVersionMixin(PublishableEntityVersionMixin):
-    """
-    Convenience mixin to link your models against ContainerVersion.
-
-    Please see docstring for ContainerVersion for more details.
-
-    If you use this class, you *MUST* also use ContainerMixin
-    """
-
-    # select these related entities by default for all queries
-    objects: ClassVar[WithRelationsManager[Self]] = WithRelationsManager(  # type: ignore[assignment]
-        "container_version",
-    )
-
-    container_version = models.OneToOneField(
-        ContainerVersion,
-        on_delete=models.CASCADE,
-    )
-
-    @property
-    def uuid(self) -> str:
-        return self.container_version.uuid
-
-    @property
-    def title(self) -> str:
-        return self.container_version.title
-
-    @property
-    def created(self) -> datetime:
-        return self.container_version.created
-
-    @property
-    def version_num(self):
-        return self.container_version.version_num
-
-    class Meta:
-        abstract = True
diff --git a/openedx_learning/apps/authoring/publishing/models/container.py b/openedx_learning/apps/authoring/publishing/models/container.py
index 4a0df08a4..0e4d99375 100644
--- a/openedx_learning/apps/authoring/publishing/models/container.py
+++ b/openedx_learning/apps/authoring/publishing/models/container.py
@@ -1,11 +1,54 @@
 """
 Container and ContainerVersion models
 """
+from typing import ClassVar, Self, TypeVar
+
+from django.core.exceptions import ValidationError
 from django.db import models
 
+from openedx_learning.lib.fields import case_sensitive_char_field
+from openedx_learning.lib.managers import WithRelationsManager
+
 from ..model_mixins.publishable_entity import PublishableEntityMixin, PublishableEntityVersionMixin
 from .entity_list import EntityList
 
+M = TypeVar('M', bound="Container")
+
+
+class ContainerManager(WithRelationsManager[M]):
+    """
+    A custom manager used for Container and its subclasses.
+    """
+    def __init__(self):
+        """
+        Initialize the manager for Container / a Container subclass
+        """
+        super().__init__(
+            # Select these related entities by default:
+            "publishable_entity",
+            "publishable_entity__published",
+            "publishable_entity__draft",
+        )
+
+    def get_queryset(self) -> models.QuerySet:
+        """
+        Apply filter() and select_related() to all querysets.
+        """
+        qs = super().get_queryset()
+        if self.model.CONTAINER_TYPE:
+            qs = qs.filter(container_type=self.model.CONTAINER_TYPE)
+        return qs
+
+    def create(self, **kwargs) -> M:
+        """
+        Apply the values from our filter when creating new instances.
+        """
+        if self.model.CONTAINER_TYPE:
+            # Don't allow creating via a subclass, like Unit.objects.create().
+            # Instead use create_container() which calls Container.objects.create(..., container_type=...)
+            raise ValidationError("Container instances should only be created via APIs like create_container()")
+        return super().create(**kwargs)
+
 
 class Container(PublishableEntityMixin):
     """
@@ -22,6 +65,40 @@ class Container(PublishableEntityMixin):
     PublishLog and Containers that were affected in a publish because their
     child elements were published.
     """
+    # Subclasses (django proxy classes) should override this
+    CONTAINER_TYPE = ""
+
+    objects: ClassVar[ContainerManager[Self]] = ContainerManager()  # type: ignore[assignment]
+
+    container_type = case_sensitive_char_field(max_length=500)
+
+    def save(self, *args, **kwargs):
+        if not self.container_type:
+            raise ValidationError("Container instances should only be created via APIs like create_container()")
+        return super().save(*args, **kwargs)
+
+    def clean(self):
+        """
+        Validate this container subclass
+        """
+        if self.container_type and self.CONTAINER_TYPE:
+            if self.container_type != self.CONTAINER_TYPE:
+                raise ValidationError("container type field mismatch with model.")
+        super().clean()
+
+    @classmethod
+    def cast_from(cls, instance: "Container") -> Self:
+        """
+        Create a new copy of a Container object, with a different subclass
+        """
+        assert instance.container_type == cls.CONTAINER_TYPE
+        new_instance = cls(
+            pk=instance.pk,
+            container_type=instance.container_type,
+        )
+        # Copy Django's internal cache of related objects
+        new_instance._state.fields_cache.update(instance._state.fields_cache)  # pylint: disable=protected-access
+        return new_instance
 
 
 class ContainerVersion(PublishableEntityVersionMixin):
@@ -58,3 +135,17 @@ class ContainerVersion(PublishableEntityVersionMixin):
         null=False,
         related_name="container_versions",
     )
+
+    @classmethod
+    def cast_from(cls, instance: "ContainerVersion") -> Self:
+        """
+        Create a new copy of a Container object, with a different subclass
+        """
+        new_instance = cls(
+            pk=instance.pk,
+            container_id=instance.container_id,
+            entity_list_id=instance.entity_list_id,
+        )
+        # Copy Django's internal cache of related objects
+        new_instance._state.fields_cache.update(instance._state.fields_cache)  # pylint: disable=protected-access
+        return new_instance
diff --git a/openedx_learning/apps/authoring/units/api.py b/openedx_learning/apps/authoring/units/api.py
index e18e6b50c..57700ff24 100644
--- a/openedx_learning/apps/authoring/units/api.py
+++ b/openedx_learning/apps/authoring/units/api.py
@@ -41,15 +41,16 @@ def create_unit(
         created: The creation date.
         created_by: The user who created the unit.
     """
-    with atomic():
-        container = publishing_api.create_container(
-            learning_package_id, key, created, created_by
-        )
-        unit = Unit.objects.create(
-            container=container,
-            publishable_entity=container.publishable_entity,
-        )
-    return unit
+    container = publishing_api.create_container(
+        learning_package_id,
+        key,
+        Unit.CONTAINER_TYPE,
+        created,
+        created_by,
+    )
+    return Unit.cast_from(container)
+    # Or, more robust but using additional queries:
+    # return Unit.objects.select_related(None).get(pk=container.pk)
 
 
 def create_unit_version(
@@ -78,22 +79,17 @@ def create_unit_version(
         created: The creation date.
         created_by: The user who created the unit.
     """
-    with atomic():
-        container_version = publishing_api.create_container_version(
-            unit.container.pk,
-            version_num,
-            title=title,
-            publishable_entities_pks=publishable_entities_pks,
-            entity_version_pks=entity_version_pks,
-            created=created,
-            created_by=created_by,
-        )
-        unit_version = UnitVersion.objects.create(
-            unit=unit,
-            container_version=container_version,
-            publishable_entity_version=container_version.publishable_entity_version,
-        )
-    return unit_version
+    container_version = publishing_api.create_container_version(
+        unit.pk,
+        version_num,
+        title=title,
+        publishable_entities_pks=publishable_entities_pks,
+        entity_version_pks=entity_version_pks,
+        created=created,
+        created_by=created_by,
+    )
+    # return UnitVersion.objects.get(pk=container_version.pk)
+    return UnitVersion.cast_from(container_version)
 
 
 def _pub_entities_for_components(
@@ -144,21 +140,15 @@ def create_next_unit_version(
         created_by: The user who created the unit.
     """
     publishable_entities_pks, entity_version_pks = _pub_entities_for_components(components)
-    with atomic():
-        container_version = publishing_api.create_next_container_version(
-            unit.container.pk,
-            title=title,
-            publishable_entities_pks=publishable_entities_pks,
-            entity_version_pks=entity_version_pks,
-            created=created,
-            created_by=created_by,
-        )
-        unit_version = UnitVersion.objects.create(
-            unit=unit,
-            container_version=container_version,
-            publishable_entity_version=container_version.publishable_entity_version,
-        )
-    return unit_version
+    container_version = publishing_api.create_next_container_version(
+        unit.pk,
+        title=title,
+        publishable_entities_pks=publishable_entities_pks,
+        entity_version_pks=entity_version_pks,
+        created=created,
+        created_by=created_by,
+    )
+    return UnitVersion.cast_from(container_version)
 
 
 def create_unit_and_version(
@@ -301,10 +291,10 @@ def get_components_in_published_unit_as_of(
     unit_pub_entity_version = publishing_api.get_published_version_as_of(unit.publishable_entity_id, publish_log_id)
     if unit_pub_entity_version is None:
         return None  # This unit was not published as of the given PublishLog ID.
-    unit_version = unit_pub_entity_version.unitversion  # type: ignore[attr-defined]
+    container_version = unit_pub_entity_version.containerversion
 
     entity_list = []
-    rows = unit_version.container_version.entity_list.entitylistrow_set.order_by("order_num")
+    rows = container_version.entity_list.entitylistrow_set.order_by("order_num")
     for row in rows:
         if row.entity_version is not None:
             component_version = row.entity_version.componentversion
diff --git a/openedx_learning/apps/authoring/units/migrations/0001_initial.py b/openedx_learning/apps/authoring/units/migrations/0001_initial.py
index 3bfc8a407..f3c8a4771 100644
--- a/openedx_learning/apps/authoring/units/migrations/0001_initial.py
+++ b/openedx_learning/apps/authoring/units/migrations/0001_initial.py
@@ -1,7 +1,6 @@
-# Generated by Django 4.2.19 on 2025-03-07 23:10
+# Generated by Django 4.2.19 on 2025-03-10 18:18
 
-import django.db.models.deletion
-from django.db import migrations, models
+from django.db import migrations
 
 
 class Migration(migrations.Migration):
@@ -16,22 +15,23 @@ class Migration(migrations.Migration):
         migrations.CreateModel(
             name='Unit',
             fields=[
-                ('publishable_entity', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, primary_key=True, serialize=False, to='oel_publishing.publishableentity')),
-                ('container', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to='oel_publishing.container')),
             ],
             options={
-                'abstract': False,
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
             },
+            bases=('oel_publishing.container',),
         ),
         migrations.CreateModel(
             name='UnitVersion',
             fields=[
-                ('publishable_entity_version', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, primary_key=True, serialize=False, to='oel_publishing.publishableentityversion')),
-                ('container_version', models.OneToOneField(on_delete=django.db.models.deletion.CASCADE, to='oel_publishing.containerversion')),
-                ('unit', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='versions', to='oel_units.unit')),
             ],
             options={
-                'abstract': False,
+                'proxy': True,
+                'indexes': [],
+                'constraints': [],
             },
+            bases=('oel_publishing.containerversion',),
         ),
     ]
diff --git a/openedx_learning/apps/authoring/units/models.py b/openedx_learning/apps/authoring/units/models.py
index 3a12e035c..60ff982c9 100644
--- a/openedx_learning/apps/authoring/units/models.py
+++ b/openedx_learning/apps/authoring/units/models.py
@@ -1,9 +1,8 @@
 """
 Models that implement units
 """
-from django.db import models
 
-from ..publishing.model_mixins import ContainerMixin, ContainerVersionMixin
+from ..publishing.models import Container, ContainerVersion
 
 __all__ = [
     "Unit",
@@ -11,21 +10,24 @@
 ]
 
 
-class Unit(ContainerMixin):
+class Unit(Container):
     """
     A Unit is Container, which is a PublishableEntity.
     """
+    CONTAINER_TYPE = "unit"
 
+    class Meta:
+        proxy = True
 
-class UnitVersion(ContainerVersionMixin):
+
+class UnitVersion(ContainerVersion):
     """
     A UnitVersion is a ContainerVersion, which is a PublishableEntityVersion.
     """
 
-    # Not sure what other metadata goes here, but we want to try to separate things
-    # like scheduling information and such into different models.
-    unit = models.ForeignKey(
-        Unit,
-        on_delete=models.CASCADE,
-        related_name="versions",
-    )
+    @property
+    def unit(self):
+        return Unit.objects.get(pk=self.container_id)
+
+    class Meta:
+        proxy = True
diff --git a/tests/openedx_learning/apps/authoring/units/test_api.py b/tests/openedx_learning/apps/authoring/units/test_api.py
index c063876d7..0a3c7ce78 100644
--- a/tests/openedx_learning/apps/authoring/units/test_api.py
+++ b/tests/openedx_learning/apps/authoring/units/test_api.py
@@ -83,22 +83,41 @@ def test_get_unit(self):
         Test get_unit()
         """
         unit = self.create_unit_with_components([self.component_1, self.component_2])
-        result = authoring_api.get_unit(unit.pk)
+        with self.assertNumQueries(1):
+            result = authoring_api.get_unit(unit.pk)
         assert result == unit
         # Versioning data should be pre-loaded via select_related()
         with self.assertNumQueries(0):
             assert result.versioning.has_unpublished_changes
-        # TODO: (maybe) This currently has extra queries and is not preloaded even though it's the same:
-        # with self.assertNumQueries(0):
-        #     assert result.container.versioning.has_unpublished_changes
+
+    def test_get_unit_non_unit(self):
+        """
+        Test that get_unit() cannot retrieve other container types
+        """
+        not_unit, _version = authoring_api.create_container_and_version(
+            self.learning_package.id,
+            key="foobar",
+            container_type="NOT a unit",  # <-- the important part
+            created=self.now,
+            created_by=None,
+            title="Testing",
+            publishable_entities_pks=[],
+            entity_version_pks=[]
+        )
+        # This generic method will work:
+        authoring_api.get_container(not_unit.pk)
+        # But the unit method will not:
+        with self.assertRaises(authoring_models.Unit.DoesNotExist):
+            authoring_api.get_unit(not_unit.pk)
 
     def test_get_container(self):
         """
         Test get_container()
         """
         unit = self.create_unit_with_components([self.component_1, self.component_2])
-        result = authoring_api.get_container(unit.container_id)
-        assert result == unit.container
+        with self.assertNumQueries(1):
+            result = authoring_api.get_container(unit.pk)
+        assert result == unit
         # Versioning data should be pre-loaded via select_related()
         with self.assertNumQueries(0):
             assert result.versioning.has_unpublished_changes
@@ -108,9 +127,9 @@ def test_create_unit_queries(self):
         Test how many database queries are required to create a unit
         """
         # The exact numbers here aren't too important - this is just to alert us if anything significant changes.
-        with self.assertNumQueries(28):
+        with self.assertNumQueries(18):
             _empty_unit = self.create_unit_with_components([])
-        with self.assertNumQueries(31):
+        with self.assertNumQueries(21):
             # And try with a non-empty unit:
             self.create_unit_with_components([self.component_1, self.component_2_v1], key="u2")
 
@@ -127,6 +146,7 @@ def test_create_unit_with_invalid_children(self):
             created=self.now,
             created_by=None,
         )
+        assert unit.versioning.draft == unit_version
         unit2, _u2v1 = authoring_api.create_unit_and_version(
             learning_package_id=self.learning_package.id,
             key="unit:key2",
@@ -144,6 +164,8 @@ def test_create_unit_with_invalid_children(self):
                 created_by=None,
             )
         # Check that a new version was not created:
+        unit.refresh_from_db()
+        assert authoring_api.get_container(unit.pk).versioning.draft == unit_version
         assert unit.versioning.draft == unit_version
 
     def test_adding_external_components(self):
@@ -272,7 +294,7 @@ def test_auto_publish_children(self):
         # Also create another component that's not in the unit at all:
         other_component, _oc_v1 = self.create_component(title="A draft component not in the unit", key="component:3")
 
-        assert authoring_api.contains_unpublished_changes(unit)
+        assert authoring_api.contains_unpublished_changes(unit.pk)
         assert self.component_1.versioning.published is None
         assert self.component_2.versioning.published is None
 
@@ -286,7 +308,7 @@ def test_auto_publish_children(self):
         self.component_1.refresh_from_db()
         assert unit.versioning.has_unpublished_changes is False  # Shallow check
         assert self.component_1.versioning.has_unpublished_changes is False
-        assert authoring_api.contains_unpublished_changes(unit) is False  # Deep check
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False  # Deep check
         assert self.component_1.versioning.published == self.component_1_v1  # v1 is now the published version.
 
         # But our other component that's outside the unit is not affected:
@@ -331,7 +353,7 @@ def test_add_component_after_publish(self):
         authoring_api.publish_all_drafts(self.learning_package.id)
         unit.refresh_from_db()  # Reloading the unit is necessary
         assert unit.versioning.has_unpublished_changes is False  # Shallow check for just the unit itself, not children
-        assert authoring_api.contains_unpublished_changes(unit) is False  # Deeper check
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False  # Deeper check
 
         # Add a published component (unpinned):
         assert self.component_1.versioning.has_unpublished_changes is False
@@ -345,7 +367,7 @@ def test_add_component_after_publish(self):
         # Now the unit should have unpublished changes:
         unit.refresh_from_db()  # Reloading the unit is necessary
         assert unit.versioning.has_unpublished_changes  # Shallow check - adding a child is a change to the unit
-        assert authoring_api.contains_unpublished_changes(unit)  # Deeper check
+        assert authoring_api.contains_unpublished_changes(unit.pk)  # Deeper check
         assert unit.versioning.draft == unit_version_v2
         assert unit.versioning.published == unit_version
 
@@ -366,7 +388,7 @@ def test_modify_unpinned_component_after_publish(self):
         unit.refresh_from_db()  # Reloading the unit is necessary if we accessed 'versioning' before publish
         self.component_1.refresh_from_db()
         assert unit.versioning.has_unpublished_changes is False  # Shallow check
-        assert authoring_api.contains_unpublished_changes(unit) is False  # Deeper check
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False  # Deeper check
         assert self.component_1.versioning.has_unpublished_changes is False
 
         # Now modify the component by changing its title (it remains a draft):
@@ -376,7 +398,7 @@ def test_modify_unpinned_component_after_publish(self):
         unit.refresh_from_db()  # Reloading the unit is necessary, or 'unit.versioning' will be outdated
         self.component_1.refresh_from_db()
         assert unit.versioning.has_unpublished_changes is False  # Shallow check should be false - unit is unchanged
-        assert authoring_api.contains_unpublished_changes(unit)  # But unit DOES contain changes
+        assert authoring_api.contains_unpublished_changes(unit.pk)  # But unit DOES contain changes
         assert self.component_1.versioning.has_unpublished_changes
 
         # Since the component changes haven't been published, they should only appear in the draft unit
@@ -395,7 +417,7 @@ def test_modify_unpinned_component_after_publish(self):
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(component_1_v2),  # new version
         ]
-        assert authoring_api.contains_unpublished_changes(unit) is False  # No longer contains unpublished changes
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False  # No longer contains unpublished changes
 
     def test_modify_pinned_component(self):
         """
@@ -420,7 +442,7 @@ def test_modify_pinned_component(self):
         unit.refresh_from_db()  # Reloading the unit is necessary, or 'unit.versioning' will be outdated
         self.component_1.refresh_from_db()
         assert unit.versioning.has_unpublished_changes is False  # Shallow check
-        assert authoring_api.contains_unpublished_changes(unit) is False  # Deep check
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False  # Deep check
         assert self.component_1.versioning.has_unpublished_changes is True
 
         # Neither the draft nor the published version of the unit is affected
@@ -496,14 +518,14 @@ def test_publishing_shared_component(self):
         unit1 = self.create_unit_with_components([c1, c2, c3], title="Unit 1", key="unit:1")
         unit2 = self.create_unit_with_components([c2, c4, c5], title="Unit 2", key="unit:2")
         authoring_api.publish_all_drafts(self.learning_package.id)
-        assert authoring_api.contains_unpublished_changes(unit1) is False
-        assert authoring_api.contains_unpublished_changes(unit2) is False
+        assert authoring_api.contains_unpublished_changes(unit1.pk) is False
+        assert authoring_api.contains_unpublished_changes(unit2.pk) is False
 
         # 2️⃣ Then the author edits C2 inside of Unit 1 making C2v2.
         c2_v2 = self.modify_component(c2, title="C2 version 2")
         # This makes U1 and U2 both show up as Units that CONTAIN unpublished changes, because they share the component.
-        assert authoring_api.contains_unpublished_changes(unit1)
-        assert authoring_api.contains_unpublished_changes(unit2)
+        assert authoring_api.contains_unpublished_changes(unit1.pk)
+        assert authoring_api.contains_unpublished_changes(unit2.pk)
         # (But the units themselves are unchanged:)
         unit1.refresh_from_db()
         unit2.refresh_from_db()
@@ -539,8 +561,8 @@ def test_publishing_shared_component(self):
         ]
 
         # Result: Unit 2 CONTAINS unpublished changes because of the modified C5. Unit 1 doesn't contain unpub changes.
-        assert authoring_api.contains_unpublished_changes(unit1) is False
-        assert authoring_api.contains_unpublished_changes(unit2)
+        assert authoring_api.contains_unpublished_changes(unit1.pk) is False
+        assert authoring_api.contains_unpublished_changes(unit2.pk)
 
         # 5️⃣ Publish component C5, which should be the only thing unpublished in the learning package
         self.publish_component(c5)
@@ -550,7 +572,7 @@ def test_publishing_shared_component(self):
             Entry(c4_v1),  # still original version of C4 (it was never modified)
             Entry(c5_v2),  # new published version of C5
         ]
-        assert authoring_api.contains_unpublished_changes(unit2) is False
+        assert authoring_api.contains_unpublished_changes(unit2.pk) is False
 
     def test_query_count_of_contains_unpublished_changes(self):
         """
@@ -570,12 +592,12 @@ def test_query_count_of_contains_unpublished_changes(self):
         authoring_api.publish_all_drafts(self.learning_package.id)
         unit.refresh_from_db()
         with self.assertNumQueries(2):
-            assert authoring_api.contains_unpublished_changes(unit) is False
+            assert authoring_api.contains_unpublished_changes(unit.pk) is False
 
         # Modify the most recently created component:
         self.modify_component(component, title="Modified Component")
         with self.assertNumQueries(2):
-            assert authoring_api.contains_unpublished_changes(unit) is True
+            assert authoring_api.contains_unpublished_changes(unit.pk) is True
 
     def test_metadata_change_doesnt_create_entity_list(self):
         """
@@ -585,14 +607,14 @@ def test_metadata_change_doesnt_create_entity_list(self):
         """
         unit = self.create_unit_with_components([self.component_1, self.component_2_v1])
 
-        orig_version_num = unit.container.versioning.draft.version_num
-        orig_entity_list_id = unit.container.versioning.draft.entity_list.pk
+        orig_version_num = unit.versioning.draft.version_num
+        orig_entity_list_id = unit.versioning.draft.entity_list.pk
 
         authoring_api.create_next_unit_version(unit, title="New Title", created=self.now)
 
         unit.refresh_from_db()
-        new_version_num = unit.container.versioning.draft.version_num
-        new_entity_list_id = unit.container.versioning.draft.entity_list.pk
+        new_version_num = unit.versioning.draft.version_num
+        new_entity_list_id = unit.versioning.draft.entity_list.pk
 
         assert new_version_num > orig_version_num
         assert new_entity_list_id == orig_entity_list_id
@@ -636,7 +658,7 @@ def test_removing_component(self):
         ]
         unit.refresh_from_db()
         assert unit.versioning.has_unpublished_changes  # The unit itself and its component list have change
-        assert authoring_api.contains_unpublished_changes(unit)
+        assert authoring_api.contains_unpublished_changes(unit.pk)
         # The published version of the unit is not yet affected:
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1),
@@ -650,7 +672,7 @@ def test_removing_component(self):
         # a footgun? We could avoid this if get_entities_in_published_container() took only an ID instead of an object,
         # but that would involve additional database lookup(s).
         unit.refresh_from_db()
-        assert authoring_api.contains_unpublished_changes(unit) is False
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1),
         ]
@@ -672,7 +694,7 @@ def test_soft_deleting_component(self):
             # reverted?
         ]
         assert unit.versioning.has_unpublished_changes is False  # The unit itself and its component list is not changed
-        assert authoring_api.contains_unpublished_changes(unit)  # But it CONTAINS an unpublished change (a deletion)
+        assert authoring_api.contains_unpublished_changes(unit.pk)  # But it CONTAINS an unpublished change (a deletion)
         # The published version of the unit is not yet affected:
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1),
@@ -681,7 +703,7 @@ def test_soft_deleting_component(self):
 
         # But when we publish the deletion, the published version is affected:
         authoring_api.publish_all_drafts(self.learning_package.id)
-        assert authoring_api.contains_unpublished_changes(unit) is False
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1),
         ]
@@ -706,7 +728,7 @@ def test_soft_deleting_and_removing_component(self):
             Entry(self.component_1_v1),
         ]
         assert unit.versioning.has_unpublished_changes is True
-        assert authoring_api.contains_unpublished_changes(unit)
+        assert authoring_api.contains_unpublished_changes(unit.pk)
         # The published version of the unit is not yet affected:
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1),
@@ -715,7 +737,7 @@ def test_soft_deleting_and_removing_component(self):
 
         # But when we publish the deletion, the published version is affected:
         authoring_api.publish_all_drafts(self.learning_package.id)
-        assert authoring_api.contains_unpublished_changes(unit) is False
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1),
         ]
@@ -734,7 +756,7 @@ def test_soft_deleting_pinned_component(self):
             Entry(self.component_2_v1, pinned=True),
         ]
         assert unit.versioning.has_unpublished_changes is False  # The unit itself and its component list is not changed
-        assert authoring_api.contains_unpublished_changes(unit) is False  # nor does it contain changes
+        assert authoring_api.contains_unpublished_changes(unit.pk) is False  # nor does it contain changes
         # The published version of the unit is also not affected:
         assert authoring_api.get_components_in_published_unit(unit) == [
             Entry(self.component_1_v1, pinned=True),
@@ -872,10 +894,7 @@ def test_units_containing(self):
         # No need to publish anything as the get_containers_with_entity() API only considers drafts (for now).
 
         with self.assertNumQueries(1):
-            result = [
-                c.unit for c in
-                authoring_api.get_containers_with_entity(self.component_1.pk).select_related("unit")
-            ]
+            result = list(authoring_api.get_containers_with_entity(self.component_1.pk))
         assert result == [
             unit1_1pinned,
             unit2_1pinned_v2,
@@ -886,10 +905,7 @@ def test_units_containing(self):
         # about pinned uses anyways (they would be unaffected by a delete).
 
         with self.assertNumQueries(1):
-            result2 = [
-                c.unit for c in
-                authoring_api.get_containers_with_entity(self.component_1.pk, ignore_pinned=True).select_related("unit")
-            ]
+            result2 = list(authoring_api.get_containers_with_entity(self.component_1.pk, ignore_pinned=True))
         assert result2 == [unit4_unpinned]
 
     # Tests TODO: