From b90c21b138f9227aaae59ec813dbf1286029b797 Mon Sep 17 00:00:00 2001
From: Nimisha Asthagiri <nasthagiri@edx.org>
Date: Tue, 20 Dec 2016 11:40:43 -0500
Subject: [PATCH] Ensure selected blocks get saved into StudentModule when
 running transform phase

TNL-6198
---
 .../transformers/library_content.py           | 27 ++++++++++++++-----
 .../tests/test_library_content.py             | 23 ++++++++--------
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/lms/djangoapps/course_blocks/transformers/library_content.py b/lms/djangoapps/course_blocks/transformers/library_content.py
index c6794471c3..c9dede0896 100644
--- a/lms/djangoapps/course_blocks/transformers/library_content.py
+++ b/lms/djangoapps/course_blocks/transformers/library_content.py
@@ -76,18 +76,34 @@ class ContentLibraryTransformer(FilteringTransformerMixin, BlockStructureTransfo
                 module = self._get_student_module(usage_info.user, usage_info.course_key, block_key)
                 if module:
                     state_dict = json.loads(module.state)
+                else:
+                    state_dict = {}
+
+                for selected_block in state_dict.get('selected', []):
                     # Add all selected entries for this user for this
                     # library module to the selected list.
-                    for state in state_dict['selected']:
-                        usage_key = usage_info.course_key.make_usage_key(state[0], state[1])
-                        if usage_key in library_children:
-                            selected.append((state[0], state[1]))
+                    block_type, block_id = selected_block
+                    usage_key = usage_info.course_key.make_usage_key(block_type, block_id)
+                    if usage_key in library_children:
+                        selected.append(selected_block)
 
-                # update selected
+                # Update selected
                 previous_count = len(selected)
                 block_keys = LibraryContentModule.make_selection(selected, library_children, max_count, mode)
                 selected = block_keys['selected']
 
+                # Save back any changes
+                if any(block_keys[changed] for changed in ('invalid', 'overlimit', 'added')):
+                    state_dict['selected'] = list(selected)
+                    StudentModule.objects.update_or_create(  # pylint: disable=no-member
+                        student=usage_info.user,
+                        course_id=usage_info.course_key,
+                        module_state_key=block_key,
+                        defaults={
+                            'state': json.dumps(state_dict),
+                        },
+                    )
+
                 # publish events for analytics
                 self._publish_events(block_structure, block_key, previous_count, max_count, block_keys)
                 all_selected_children.update(usage_info.course_key.make_usage_key(s[0], s[1]) for s in selected)
@@ -125,7 +141,6 @@ class ContentLibraryTransformer(FilteringTransformerMixin, BlockStructureTransfo
                 student=user,
                 course_id=course_key,
                 module_state_key=block_key,
-                state__contains='"selected": [['
             )
         except StudentModule.DoesNotExist:
             return None
diff --git a/lms/djangoapps/course_blocks/transformers/tests/test_library_content.py b/lms/djangoapps/course_blocks/transformers/tests/test_library_content.py
index 691597410e..846081b69f 100644
--- a/lms/djangoapps/course_blocks/transformers/tests/test_library_content.py
+++ b/lms/djangoapps/course_blocks/transformers/tests/test_library_content.py
@@ -1,7 +1,7 @@
 """
 Tests for ContentLibraryTransformer.
 """
-import mock
+
 from student.tests.factories import CourseEnrollmentFactory
 
 from openedx.core.djangoapps.content.block_structure.api import clear_course_from_cache
@@ -138,14 +138,14 @@ class ContentLibraryTransformerTestCase(CourseStructureTestCase):
 
         vertical2_selected = self.get_block_key_set(self.blocks, 'vertical2').pop() in trans_keys
         vertical3_selected = self.get_block_key_set(self.blocks, 'vertical3').pop() in trans_keys
-        self.assertTrue(vertical2_selected or vertical3_selected)
 
-        # Check course structure again, with mocked selected modules for a user.
-        with mock.patch(
-            'lms.djangoapps.course_blocks.transformers.library_content.ContentLibraryTransformer._get_student_module',
-            return_value=MockedModule('{"selected": [["vertical", "vertical_vertical2"]]}'),
-        ):
-            clear_course_from_cache(self.course.id)
+        self.assertNotEquals(vertical2_selected, vertical3_selected)  # only one of them should be selected
+        selected_vertical = 'vertical2' if vertical2_selected else 'vertical3'
+        selected_child = 'html1' if vertical2_selected else 'html2'
+
+        # Check course structure again.
+        clear_course_from_cache(self.course.id)
+        for i in range(5):
             trans_block_structure = get_course_blocks(
                 self.user,
                 self.course.location,
@@ -160,7 +160,8 @@ class ContentLibraryTransformerTestCase(CourseStructureTestCase):
                     'lesson1',
                     'vertical1',
                     'library_content1',
-                    'vertical2',
-                    'html1'
-                )
+                    selected_vertical,
+                    selected_child,
+                ),
+                "Expected 'selected' equality failed in iteration {}.".format(i)
             )