1818 LearningPackage ,
1919 PublishableEntity ,
2020)
21+ from openedx_learning .apps .authoring .units .models import Unit
2122from openedx_learning .lib .test_utils import TestCase
2223
2324User = get_user_model ()
@@ -221,10 +222,11 @@ def test_create_collection_without_description(self):
221222
222223class CollectionEntitiesTestCase (CollectionsTestCase ):
223224 """
224- Base class with collections that contain components .
225+ Base class with collections that contain entities .
225226 """
226227 published_component : Component
227228 draft_component : Component
229+ draft_unit : Unit
228230 user : UserType
229231 html_type : ComponentType
230232 problem_type : ComponentType
@@ -243,6 +245,13 @@ def setUpTestData(cls) -> None:
243245
244246 cls .html_type = api .get_or_create_component_type ("xblock.v1" , "html" )
245247 cls .problem_type = api .get_or_create_component_type ("xblock.v1" , "problem" )
248+ created_time = datetime (2025 , 4 , 1 , tzinfo = timezone .utc )
249+ cls .draft_unit = api .create_unit (
250+ learning_package_id = cls .learning_package .id ,
251+ key = "unit-1" ,
252+ created = created_time ,
253+ created_by = cls .user .id ,
254+ )
246255
247256 # Make and publish one Component
248257 cls .published_component , _ = api .create_component_and_version (
@@ -284,6 +293,7 @@ def setUpTestData(cls) -> None:
284293 entities_qset = PublishableEntity .objects .filter (id__in = [
285294 cls .published_component .pk ,
286295 cls .draft_component .pk ,
296+ cls .draft_unit .pk ,
287297 ]),
288298 )
289299 cls .disabled_collection = api .add_to_collection (
@@ -308,6 +318,7 @@ def test_create_collection_entities(self):
308318 self .published_component .publishable_entity ,
309319 ]
310320 assert list (self .collection2 .entities .all ()) == [
321+ self .draft_unit .publishable_entity ,
311322 self .published_component .publishable_entity ,
312323 self .draft_component .publishable_entity ,
313324 ]
@@ -325,11 +336,13 @@ def test_add_to_collection(self):
325336 self .collection1 .key ,
326337 PublishableEntity .objects .filter (id__in = [
327338 self .draft_component .pk ,
339+ self .draft_unit .pk ,
328340 ]),
329341 created_by = self .user .id ,
330342 )
331343
332344 assert list (self .collection1 .entities .all ()) == [
345+ self .draft_unit .publishable_entity ,
333346 self .published_component .publishable_entity ,
334347 self .draft_component .publishable_entity ,
335348 ]
@@ -352,6 +365,7 @@ def test_add_to_collection_again(self):
352365 )
353366
354367 assert list (self .collection2 .entities .all ()) == [
368+ self .draft_unit .publishable_entity ,
355369 self .published_component .publishable_entity ,
356370 self .draft_component .publishable_entity ,
357371 ]
@@ -383,6 +397,7 @@ def test_remove_from_collection(self):
383397 self .collection2 .key ,
384398 PublishableEntity .objects .filter (id__in = [
385399 self .published_component .pk ,
400+ self .draft_unit .pk ,
386401 ]),
387402 )
388403
@@ -422,6 +437,24 @@ def test_get_collection_components(self):
422437 self .another_library_collection .key ,
423438 ))
424439
440+ def test_get_collection_containers (self ):
441+ assert not list (api .get_collection_containers (
442+ self .learning_package .id ,
443+ self .collection1 .key ,
444+ ))
445+ assert list (api .get_collection_containers (
446+ self .learning_package .id ,
447+ self .collection2 .key ,
448+ )) == [self .draft_unit .container ]
449+ assert not list (api .get_collection_containers (
450+ self .learning_package .id ,
451+ self .collection3 .key ,
452+ ))
453+ assert not list (api .get_collection_containers (
454+ self .learning_package .id ,
455+ self .another_library_collection .key ,
456+ ))
457+
425458
426459class UpdateCollectionTestCase (CollectionTestCase ):
427460 """
@@ -533,6 +566,7 @@ def test_soft_delete(self):
533566 assert collection == api .get_collection (self .learning_package .id , collection .key )
534567 # ...and the collection's entities remain intact.
535568 assert list (collection .entities .all ()) == [
569+ self .draft_unit .publishable_entity ,
536570 self .published_component .publishable_entity ,
537571 self .draft_component .publishable_entity ,
538572 ]
@@ -586,6 +620,7 @@ def test_restore(self):
586620 assert collection == api .get_collection (self .learning_package .id , collection .key )
587621 # ...and the collection's entities remain intact.
588622 assert list (collection .entities .all ()) == [
623+ self .draft_unit .publishable_entity ,
589624 self .published_component .publishable_entity ,
590625 self .draft_component .publishable_entity ,
591626 ]
@@ -615,6 +650,7 @@ def test_set_collections(self):
615650 self .draft_component .publishable_entity ,
616651 ]
617652 assert list (self .collection2 .entities .all ()) == [
653+ self .draft_unit .publishable_entity ,
618654 self .published_component .publishable_entity ,
619655 self .draft_component .publishable_entity ,
620656 ]
@@ -651,6 +687,7 @@ def test_set_collections(self):
651687 self .published_component .publishable_entity ,
652688 ]
653689 assert list (self .collection2 .entities .all ()) == [
690+ self .draft_unit .publishable_entity ,
654691 self .published_component .publishable_entity ,
655692 self .draft_component .publishable_entity ,
656693 ]
0 commit comments