diff --git a/cms/djangoapps/contentstore/features/component_settings_editor_helpers.py b/cms/djangoapps/contentstore/features/component_settings_editor_helpers.py
index 8c6a1023c9..a130a33677 100644
--- a/cms/djangoapps/contentstore/features/component_settings_editor_helpers.py
+++ b/cms/djangoapps/contentstore/features/component_settings_editor_helpers.py
@@ -12,7 +12,7 @@ def create_component_instance(step, component_button_css, category,
                               has_multiple_templates=True):
 
     click_new_component_button(step, component_button_css)
-    if category == 'problem' or category == 'html':
+    if category in ('problem', 'html'):
         def animation_done(_driver):
             return world.browser.evaluate_script("$('div.new-component').css('display')") == 'none'
         world.wait_for(animation_done)
diff --git a/cms/djangoapps/contentstore/features/problem-editor.py b/cms/djangoapps/contentstore/features/problem-editor.py
index e97e2857a5..d7ccb557ba 100644
--- a/cms/djangoapps/contentstore/features/problem-editor.py
+++ b/cms/djangoapps/contentstore/features/problem-editor.py
@@ -157,7 +157,7 @@ def create_latex_problem(step):
     world.click_new_component_button(step, '.large-problem-icon')
 
     def animation_done(_driver):
-            return world.browser.evaluate_script("$('div.new-component').css('display')") == 'none'
+        return world.browser.evaluate_script("$('div.new-component').css('display')") == 'none'
     world.wait_for(animation_done)
     # Go to advanced tab.
     world.css_click('#ui-id-2')