LMS-11357 Split revert_to_publish: correctly revert to Published version

common/lib/xmodule/xmodule/modulestore/split_mongo/split.py

@@ -363,6 +363,9 @@ class BulkWriteMixin(object):
         """
         Copy the structure and update the history info (edited_by, edited_on, previous_version)
         """
+        if course_key.branch is None:
+            raise InsufficientSpecificationError(course_key)
+
         bulk_write_record = self._get_bulk_write_record(course_key)
 
         # If we have an active bulk write, and it's already been edited, then just use that structure
@@ -1891,16 +1894,7 @@ class SplitMongoModuleStore(BulkWriteMixin, ModuleStoreWriteBase):
                 parent_block['edit_info']['previous_version'] = parent_block['edit_info']['update_version']
                 parent_block['edit_info']['update_version'] = new_id
 
-            def remove_subtree(block_id):
-                """
-                Remove the subtree rooted at block_id
-                """
-                encoded_block_id = encode_key_for_mongo(block_id)
-                for child in new_blocks[encoded_block_id]['fields'].get('children', []):
-                    remove_subtree(child)
-                del new_blocks[encoded_block_id]
-
-            remove_subtree(usage_locator.block_id)
+            self._remove_subtree(usage_locator.block_id, new_blocks)
 
             # update index if appropriate and structures
             self.update_structure(usage_locator.course_key, new_structure)
@@ -1914,6 +1908,15 @@ class SplitMongoModuleStore(BulkWriteMixin, ModuleStoreWriteBase):
 
             return result
 
+    def _remove_subtree(self, block_id, blocks):
+        """
+        Remove the subtree rooted at block_id
+        """
+        encoded_block_id = encode_key_for_mongo(block_id)
+        for child in blocks[encoded_block_id]['fields'].get('children', []):
+            self._remove_subtree(child, blocks)
+        del blocks[encoded_block_id]
+
     def delete_course(self, course_key, user_id):
         """
         Remove the given course from the course index.

common/lib/xmodule/xmodule/modulestore/split_mongo/split_draft.py

@@ -307,19 +307,45 @@ class DraftVersioningModuleStore(ModuleStoreDraftAndPublished, SplitMongoModuleS
         if location.category in DIRECT_ONLY_CATEGORIES:
             return
 
-        if not self.has_item(location, revision=ModuleStoreEnum.RevisionOption.published_only):
-            raise InvalidVersionError(location)
+        draft_course_key = location.course_key.for_branch(ModuleStoreEnum.BranchName.draft)
+        with self.bulk_operations(draft_course_key):
 
-        SplitMongoModuleStore.copy(
-            self,
-            user_id,
-            # Directly using the replace function rather than the for_branch function
-            # because for_branch obliterates the version_guid and will lead to missed version conflicts.
-            # TODO Instead, the for_branch implementation should be fixed in the Opaque Keys library.
-            location.course_key.replace(branch=ModuleStoreEnum.BranchName.published),
-            location.course_key.for_branch(ModuleStoreEnum.BranchName.draft),
-            [location]
-        )
+            # get head version of Published branch
+            published_course_structure = self._lookup_course(
+                location.course_key.for_branch(ModuleStoreEnum.BranchName.published)
+            )['structure']
+            published_block = self._get_block_from_structure(published_course_structure, location.block_id)
+            if published_block is None:
+                raise InvalidVersionError(location)
+
+            # create a new versioned draft structure
+            draft_course_structure = self._lookup_course(draft_course_key)['structure']
+            new_structure = self.version_structure(draft_course_key, draft_course_structure, user_id)
+
+            # remove the block and its descendants from the new structure
+            self._remove_subtree(location.block_id, new_structure['blocks'])
+
+            # copy over the block and its descendants from the published branch
+            def copy_from_published(root_block_id):
+                """
+                copies root_block_id and its descendants from published_course_structure to new_structure
+                """
+                self._update_block_in_structure(
+                    new_structure,
+                    root_block_id,
+                    self._get_block_from_structure(published_course_structure, root_block_id)
+                )
+                block = self._get_block_from_structure(new_structure, root_block_id)
+                for child_block_id in block.setdefault('fields', {}).get('children', []):
+                    copy_from_published(child_block_id)
+
+            copy_from_published(location.block_id)
+
+            # update course structure and index
+            self.update_structure(draft_course_key, new_structure)
+            index_entry = self._get_index_if_valid(draft_course_key)
+            if index_entry is not None:
+                self._update_head(draft_course_key, index_entry, ModuleStoreEnum.BranchName.draft, new_structure['_id'])
 
     def get_course_history_info(self, course_locator):
         """

common/lib/xmodule/xmodule/modulestore/tests/test_mixed_modulestore.py

@@ -982,9 +982,11 @@ class TestMixedModuleStore(unittest.TestCase):
         vertical_children_num = len(vertical.children)
 
         self.store.publish(self.course.location, self.user_id)
+        self.assertFalse(self._has_changes(self.vertical_x1a))
 
         # delete leaf problem (will make parent vertical a draft)
         self.store.delete_item(self.problem_x1a_1, self.user_id)
+        self.assertTrue(self._has_changes(self.vertical_x1a))
 
         draft_parent = self.store.get_item(self.vertical_x1a)
         self.assertEqual(vertical_children_num - 1, len(draft_parent.children))
@@ -998,6 +1000,7 @@ class TestMixedModuleStore(unittest.TestCase):
         reverted_parent = self.store.get_item(self.vertical_x1a)
         self.assertEqual(vertical_children_num, len(published_parent.children))
         self.assertEqual(reverted_parent, published_parent)
+        self.assertFalse(self._has_changes(self.vertical_x1a))
 
     @ddt.data('draft', 'split')
     def test_revert_to_published_root_published(self, default_ms):

common/lib/xmodule/xmodule/modulestore/tests/test_split_modulestore_bulk_operations.py

@@ -18,8 +18,8 @@ class TestBulkWriteMixin(unittest.TestCase):
         self.conn = self.bulk.db_connection = MagicMock(name='db_connection', spec=MongoConnection)
         self.conn.get_course_index.return_value = {'initial': 'index'}
 
-        self.course_key = CourseLocator('org', 'course', 'run-a')
-        self.course_key_b = CourseLocator('org', 'course', 'run-b')
+        self.course_key = CourseLocator('org', 'course', 'run-a', branch='test')
+        self.course_key_b = CourseLocator('org', 'course', 'run-b', branch='test')
         self.structure = {'this': 'is', 'a': 'structure', '_id': ObjectId()}
         self.index_entry = {'this': 'is', 'an': 'index'}