feat: component set collection api [FC-0062]

openedx_learning/apps/authoring/components/api.py

@@ -12,16 +12,18 @@
 """
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import StrEnum, auto
 from logging import getLogger
 from pathlib import Path
 from uuid import UUID
 
+from django.core.exceptions import ValidationError
 from django.db.models import Q, QuerySet
 from django.db.transaction import atomic
 from django.http.response import HttpResponse, HttpResponseNotFound
 
+from ..collections.models import Collection, CollectionPublishableEntity
 from ..contents import api as contents_api
 from ..publishing import api as publishing_api
 from .models import Component, ComponentType, ComponentVersion, ComponentVersionContent
@@ -48,6 +50,7 @@
     "look_up_component_version_content",
     "AssetError",
     "get_redirect_response_for_component_asset",
+    "set_collections",
 ]
 
 
@@ -603,3 +606,54 @@ def _error_header(error: AssetError) -> dict[str, str]:
     )
 
     return HttpResponse(headers={**info_headers, **redirect_headers})
+
+
+def set_collections(
+    learning_package_id: int,
+    component: Component,
+    collection_qset: QuerySet[Collection],
+    created_by: int | None = None,
+) -> set[Collection]:
+    """
+    Set collections for a given component.
+
+    These Collections must belong to the same LearningPackage as the Component, or a ValidationError will be raised.
+
+    Modified date of all collections related to component is updated.
+
+    Returns the updated collections.
+    """
+    # Disallow adding entities outside the collection's learning package
+    invalid_collection = collection_qset.exclude(learning_package_id=learning_package_id).first()
+    if invalid_collection:
+        raise ValidationError(
+            f"Cannot add collection {invalid_collection.pk} in learning package  "
+            f"{invalid_collection.learning_package_id} to component {component} in "
+            f"learning package {learning_package_id}."
+        )
+    current_relations = CollectionPublishableEntity.objects.filter(
+        entity=component.publishable_entity
+    ).select_related('collection')
+    # Clear other collections for given component and add only new collections from collection_qset
+    removed_collections = set(
+        r.collection for r in current_relations.exclude(collection__in=collection_qset)
+    )
+    new_collections = set(collection_qset.exclude(
+        id__in=current_relations.values_list('collection', flat=True)
+    ))
+    # Use `remove` instead of `CollectionPublishableEntity.delete()` to trigger m2m_changed signal which will handle
+    # updating component index.
+    component.publishable_entity.collections.remove(*removed_collections)
+    component.publishable_entity.collections.add(
+        *new_collections,
+        through_defaults={"created_by_id": created_by},
+    )
+    # Update modified date via update to avoid triggering post_save signal for collections
+    # The signal triggers index update for each collection synchronously which will be very slow in this case.
+    # Instead trigger the index update in the caller function asynchronously.
+    affected_collection = removed_collections | new_collections
+    Collection.objects.filter(
+        id__in=[collection.id for collection in affected_collection]
+    ).update(modified=datetime.now(tz=timezone.utc))
+
+    return affected_collection

tests/openedx_learning/apps/authoring/components/test_api.py

@@ -3,8 +3,12 @@
 """
 from datetime import datetime, timezone
 
-from django.core.exceptions import ObjectDoesNotExist
+from django.contrib.auth import get_user_model
+from django.core.exceptions import ObjectDoesNotExist, ValidationError
+from freezegun import freeze_time
 
+from openedx_learning.apps.authoring.collections import api as collection_api
+from openedx_learning.apps.authoring.collections.models import Collection, CollectionPublishableEntity
 from openedx_learning.apps.authoring.components import api as components_api
 from openedx_learning.apps.authoring.components.models import Component, ComponentType
 from openedx_learning.apps.authoring.contents import api as contents_api
@@ -13,6 +17,8 @@
 from openedx_learning.apps.authoring.publishing.models import LearningPackage
 from openedx_learning.lib.test_utils import TestCase
 
+User = get_user_model()
+
 
 class ComponentTestCase(TestCase):
     """
@@ -503,3 +509,128 @@ def test_multiple_versions(self):
             version_3.contents
                      .get(componentversioncontent__key="hello.txt")
         )
+
+
+class SetCollectionsTestCase(ComponentTestCase):
+    """
+    Test setting collections for a component.
+    """
+    collection1: Collection
+    collection2: Collection
+    collection3: Collection
+    published_problem: Component
+    user: User  # type: ignore [valid-type]
+
+    @classmethod
+    def setUpTestData(cls) -> None:
+        """
+        Initialize some collections
+        """
+        super().setUpTestData()
+        v2_problem_type = components_api.get_or_create_component_type("xblock.v2", "problem")
+        cls.published_problem, _ = components_api.create_component_and_version(
+            cls.learning_package.id,
+            component_type=v2_problem_type,
+            local_key="pp_lk",
+            title="Published Problem",
+            created=cls.now,
+            created_by=None,
+        )
+        cls.collection1 = collection_api.create_collection(
+            cls.learning_package.id,
+            key="MYCOL1",
+            title="Collection1",
+            created_by=None,
+            description="Description of Collection 1",
+        )
+        cls.collection2 = collection_api.create_collection(
+            cls.learning_package.id,
+            key="MYCOL2",
+            title="Collection2",
+            created_by=None,
+            description="Description of Collection 2",
+        )
+        cls.collection3 = collection_api.create_collection(
+            cls.learning_package.id,
+            key="MYCOL3",
+            title="Collection3",
+            created_by=None,
+            description="Description of Collection 3",
+        )
+        cls.user = User.objects.create(
+            username="user",
+            email="[email protected]",
+        )
+
+    def test_set_collections(self):
+        """
+        Test setting collections in a component
+        """
+        modified_time = datetime(2024, 8, 8, tzinfo=timezone.utc)
+        with freeze_time(modified_time):
+            components_api.set_collections(
+                self.learning_package.id,
+                self.published_problem,
+                collection_qset=Collection.objects.filter(id__in=[
+                    self.collection1.pk,
+                    self.collection2.pk,
+                ]),
+                created_by=self.user.id,
+            )
+        assert list(self.collection1.entities.all()) == [
+            self.published_problem.publishable_entity,
+        ]
+        assert list(self.collection2.entities.all()) == [
+            self.published_problem.publishable_entity,
+        ]
+        for collection_entity in CollectionPublishableEntity.objects.filter(
+            entity=self.published_problem.publishable_entity
+        ):
+            assert collection_entity.created_by == self.user
+        assert Collection.objects.get(id=self.collection1.pk).modified == modified_time
+        assert Collection.objects.get(id=self.collection2.pk).modified == modified_time
+
+        # Set collections again, but this time remove collection1 and add collection3
+        # Expected result: collection2 & collection3 associated to component and collection1 is excluded.
+        new_modified_time = datetime(2024, 8, 8, tzinfo=timezone.utc)
+        with freeze_time(new_modified_time):
+            components_api.set_collections(
+                self.learning_package.id,
+                self.published_problem,
+                collection_qset=Collection.objects.filter(id__in=[
+                    self.collection3.pk,
+                    self.collection2.pk,
+                ]),
+                created_by=self.user.id,
+            )
+        assert not list(self.collection1.entities.all())
+        assert list(self.collection2.entities.all()) == [
+            self.published_problem.publishable_entity,
+        ]
+        assert list(self.collection3.entities.all()) == [
+            self.published_problem.publishable_entity,
+        ]
+        # update modified time of all three collections as they were all updated
+        assert Collection.objects.get(id=self.collection1.pk).modified == new_modified_time
+        assert Collection.objects.get(id=self.collection2.pk).modified == new_modified_time
+        assert Collection.objects.get(id=self.collection3.pk).modified == new_modified_time
+
+    def test_set_collection_wrong_learning_package(self):
+        """
+        We cannot set collections with a different learning package than the component.
+        """
+        learning_package_2 = publishing_api.create_learning_package(
+            key="ComponentTestCase-test-key-2",
+            title="Components Test Case Learning Package-2",
+        )
+        with self.assertRaises(ValidationError):
+            components_api.set_collections(
+                learning_package_2.id,
+                self.published_problem,
+                collection_qset=Collection.objects.filter(id__in=[
+                    self.collection1.pk,
+                ]),
+                created_by=self.user.id,
+            )
+
+        assert not list(self.collection1.entities.all())