Update factory definitions

server/books/blocks.py

@@ -76,7 +76,9 @@ class AssignmentBlock(blocks.StructBlock):
     class Meta:
         icon = "download"
 
-    assignment_id = SnippetChooserBlock(Assignment)
+    assignment_id = SnippetChooserBlock(
+        Assignment
+    )  # probably wrong name, should be `assignment` only. todo: Can we change it?
 
 
 # 'survey'
@@ -84,7 +86,9 @@ class SurveyBlock(blocks.StructBlock):
     class Meta:
         icon = "form"
 
-    survey_id = SnippetChooserBlock(Survey)
+    survey_id = SnippetChooserBlock(
+        Survey
+    )  # probably wrong name, should be `survey` only. todo: Can we change it?
 
 
 # 'video_block'

server/books/factories.py

@@ -176,33 +176,15 @@ class LinkBlockFactory(wagtail_factories.StructBlockFactory):
         model = LinkBlock
 
 
-class EntityBlockFactory(wagtail_factories.StructBlockFactory):
-    @classmethod
-    def _build(cls, model_class, *args, **kwargs):
-        block = model_class()
-        logger.debug(cls.id_key)
-        logger.debug(cls.entity_key)
-        logger.debug(kwargs)
-        value = block.to_python({cls.id_key: kwargs.get(cls.entity_key).id})
-        clean_value = block.clean(value)
-        return clean_value
-
-
-class AssignmentBlockFactory(EntityBlockFactory):
+class AssignmentBlockFactory(wagtail_factories.StructBlockFactory):
     class Meta:
         model = AssignmentBlock
 
-    id_key = "assignment_id"
-    entity_key = "assignment"
 
-
-class SurveyBlockFactory(EntityBlockFactory):
+class SurveyBlockFactory(wagtail_factories.StructBlockFactory):
     class Meta:
         model = SurveyBlock
 
-    id_key = "survey_id"
-    entity_key = "survey"
-
 
 class VideoBlockFactory(wagtail_factories.StructBlockFactory):
     url = factory.LazyAttribute(lambda x: "https://www.youtube.com/watch?v=lO9d-AJai8Q")