diff --git a/openassessment/xblock/ui_mixins/mfe/assessment_serializers.py b/openassessment/xblock/ui_mixins/mfe/assessment_serializers.py
index 1ad6c86419..54822e67ef 100644
--- a/openassessment/xblock/ui_mixins/mfe/assessment_serializers.py
+++ b/openassessment/xblock/ui_mixins/mfe/assessment_serializers.py
@@ -3,13 +3,14 @@
 """
 # pylint: disable=abstract-method
 
-from rest_framework.fields import (
+from rest_framework.serializers import (
     CharField,
     IntegerField,
     SerializerMethodField,
     URLField,
     Serializer,
 )
+from openassessment.xblock.ui_mixins.mfe.serializer_utils import NullField
 
 
 class AssessmentScoreSerializer(Serializer):
@@ -56,8 +57,6 @@ class AssessmentStepSerializer(Serializer):
     stepScore = AssessmentScoreSerializer(source="*")
     assessment = AssessmentDataSerializer(source="*")
 
-from openassessment.xblock.ui_mixins.mfe.serializer_utils import NullField
-
 
 class SubmissionFileSerializer(Serializer):
     fileUrl = URLField(source="file_key")
diff --git a/openassessment/xblock/ui_mixins/mfe/test_assessment_serializers.py b/openassessment/xblock/ui_mixins/mfe/test_assessment_serializers.py
index 79767ae278..10fe240630 100644
--- a/openassessment/xblock/ui_mixins/mfe/test_assessment_serializers.py
+++ b/openassessment/xblock/ui_mixins/mfe/test_assessment_serializers.py
@@ -25,7 +25,7 @@ class TestAssessmentResponseSerializer(XBlockHandlerTestCase, SubmissionTestMixi
     @scenario("data/basic_scenario.xml", user_id="Alan")
     def test_no_response(self, xblock):
         # Given we are asking for assessment data too early (still on submission step)
-        context = {"response": None}
+        context = {"response": None, "step": "submission"}
 
         # When I load my response
         data = AssessmentResponseSerializer(xblock.api_data, context=context).data
@@ -47,7 +47,7 @@ def test_response(self, xblock):
         submission = self.create_test_submission(
             xblock, submission_text=submission_text
         )
-        context = {"response": submission}
+        context = {"response": submission, "step": "self"}
 
         # When I load my response
         data = AssessmentResponseSerializer(xblock.api_data, context=context).data
@@ -72,7 +72,7 @@ def test_files_empty(self, xblock):
         submission = self.create_test_submission(
             xblock, submission_text=submission_text
         )
-        context = {"response": submission}
+        context = {"response": submission, "step": "self"}
 
         # When I load my response
         data = AssessmentResponseSerializer(xblock.api_data, context=context).data
@@ -138,7 +138,7 @@ def test_files(self, xblock, mock_get_files):
         )
 
         # When I load my response
-        context = {"response": submission}
+        context = {"response": submission, "step": "self"}
         data = AssessmentResponseSerializer(xblock.api_data, context=context).data
 
         # I get the appropriate response (test URLs use usage ID)