diff --git a/lms/djangoapps/courseware/tests/test_word_cloud.py b/lms/djangoapps/courseware/tests/test_word_cloud.py
index f36b847847a..2e6d939e5d0 100644
--- a/lms/djangoapps/courseware/tests/test_word_cloud.py
+++ b/lms/djangoapps/courseware/tests/test_word_cloud.py
@@ -6,6 +6,8 @@
 import json
 from operator import itemgetter
 
+from django.conf import settings
+
 # noinspection PyUnresolvedReferences
 from xmodule.tests.helpers import override_descriptor_system, mock_render_template  # pylint: disable=unused-import
 from xmodule.x_module import STUDENT_VIEW
@@ -226,9 +228,13 @@ def test_word_cloud_constructor(self, mock_render_django_template):
             'instructions': self.block.instructions,
             'element_id': self.block.location.html_id(),
             'num_inputs': 5,  # default value
-            'range_num_inputs': range(5),
             'submitted': False,  # default value,
         }
 
-        mock_render_django_template.assert_called_once()
-        assert fragment.content == self.runtime.render_template('templates/word_cloud.html', expected_context)
+        if settings.USE_EXTRACTED_WORD_CLOUD_BLOCK:
+            expected_context['range_num_inputs'] = range(5)
+            mock_render_django_template.assert_called_once()
+            assert fragment.content == self.runtime.render_template('templates/word_cloud.html', expected_context)
+        else:
+            expected_context['element_class'] = self.block.location.block_type
+            assert fragment.content == self.runtime.render_template('word_cloud.html', expected_context)
diff --git a/xmodule/tests/test_word_cloud.py b/xmodule/tests/test_word_cloud.py
index 29b20995ea3..e5d05c3ce28 100644
--- a/xmodule/tests/test_word_cloud.py
+++ b/xmodule/tests/test_word_cloud.py
@@ -4,6 +4,7 @@
 import os
 from unittest.mock import Mock
 
+from django.conf import settings
 from django.test import TestCase
 from fs.memoryfs import MemoryFS
 from lxml import etree
@@ -38,15 +39,20 @@ def test_xml_import_export_cycle(self):
         runtime.export_fs = MemoryFS()
 
         original_xml = (
-            '<word_cloud xblock-family="xblock.v1" display_name="Favorite Fruits" display_student_percents="false" '
+            '<word_cloud display_name="Favorite Fruits" display_student_percents="false" '
             'instructions="What are your favorite fruits?" num_inputs="3" num_top_words="100"/>\n'
         )
 
         olx_element = etree.fromstring(original_xml)
         runtime.id_generator = Mock()
-        def_id = runtime.id_generator.create_definition(olx_element.tag, olx_element.get('url_name'))
-        keys = ScopeIds(None, olx_element.tag, def_id, runtime.id_generator.create_usage(def_id))
-        block = WordCloudBlock.parse_xml(olx_element, runtime, keys)
+
+        if settings.USE_EXTRACTED_WORD_CLOUD_BLOCK:
+            def_id = runtime.id_generator.create_definition(olx_element.tag, olx_element.get('url_name'))
+            keys = ScopeIds(None, olx_element.tag, def_id, runtime.id_generator.create_usage(def_id))
+            block = WordCloudBlock.parse_xml(olx_element, runtime, keys)
+        else:
+            block = WordCloudBlock.parse_xml(olx_element, runtime, None)
+
         block.location = BlockUsageLocator(
             CourseLocator('org', 'course', 'run', branch='revision'), 'word_cloud', 'block_id'
         )
@@ -57,19 +63,27 @@ def test_xml_import_export_cycle(self):
         assert block.num_inputs == 3
         assert block.num_top_words == 100
 
-        filepath = 'word_cloud/block_id.xml'
-        runtime.export_fs.makedirs(os.path.dirname(filepath), recreate=True)
-        with runtime.export_fs.open(filepath, 'wb') as fileobj:
-            runtime.export_to_xml(block, fileobj)
+        if settings.USE_EXTRACTED_WORD_CLOUD_BLOCK:
+            filepath = 'word_cloud/block_id.xml'
+            runtime.export_fs.makedirs(os.path.dirname(filepath), recreate=True)
+            with runtime.export_fs.open(filepath, 'wb') as fileObj:
+                runtime.export_to_xml(block, fileObj)
+        else:
+            node = etree.Element("unknown_root")
+            # This will export the olx to a separate file.
+            block.add_xml_to_node(node)
 
         with runtime.export_fs.open('word_cloud/block_id.xml') as f:
             exported_xml = f.read()
 
-        exported_xml_tree = etree.fromstring(exported_xml.encode('utf-8'))
-        etree.cleanup_namespaces(exported_xml_tree)
-        exported_xml = etree.tostring(exported_xml_tree, encoding='unicode', pretty_print=True)
+        if settings.USE_EXTRACTED_WORD_CLOUD_BLOCK:
+            exported_xml_tree = etree.fromstring(exported_xml.encode('utf-8'))
+            etree.cleanup_namespaces(exported_xml_tree)
+            if 'xblock-family' in exported_xml_tree.attrib:
+                del exported_xml_tree.attrib['xblock-family']
+            exported_xml = etree.tostring(exported_xml_tree, encoding='unicode', pretty_print=True)
 
-        assert original_xml == exported_xml
+        assert exported_xml == original_xml
 
     def test_bad_ajax_request(self):
         """