diff --git a/apps/notification/schema.py b/apps/notification/schema.py
index 961f9e0fc1..2302bb9e14 100644
--- a/apps/notification/schema.py
+++ b/apps/notification/schema.py
@@ -76,6 +76,7 @@ class Meta:
     project = graphene.Field(AssignmentProjectDetailType)
     content_data = graphene.Field(AssignmentContentDataType)
 
+    @staticmethod
     def resolve_content_data(root, info):
         return info.context.dl.notification.assignment.load(root.pk)
 
diff --git a/apps/notification/tests/test_apis.py b/apps/notification/tests/test_apis.py
index 9c6d2498a2..0a5c9ddc2a 100644
--- a/apps/notification/tests/test_apis.py
+++ b/apps/notification/tests/test_apis.py
@@ -506,13 +506,15 @@ def test_assignment_create_on_entry_comment_assignee_change(self):
         assert data['count'] == 1  # assignment for user2
 
     def test_assignment_is_done(self):
+        # XXX: To avoid using content type cache from pre-tests
+        ContentType.objects.clear_cache()
+
         project = self.create(Project)
         user1 = self.create(User)
         user2 = self.create(User)
         lead = self.create(Lead, project=project)
         kwargs = {
-            'object_id': lead.id,
-            'content_type': ContentType.objects.get_for_model(Lead),
+            'content_object': lead,
             'project': project,
             'created_for': user1,
             'created_by': user2,
diff --git a/apps/notification/tests/test_schemas.py b/apps/notification/tests/test_schemas.py
index afc2571313..eecf3b1bdb 100644
--- a/apps/notification/tests/test_schemas.py
+++ b/apps/notification/tests/test_schemas.py
@@ -2,16 +2,16 @@
 import pytz
 
 from django.contrib.contenttypes.models import ContentType
+
 from utils.graphene.tests import GraphQLTestCase
-from lead.models import Lead
-from notification.models import Assignment, Notification
-from quality_assurance.models import EntryReviewComment
+from notification.models import Notification
 
 from user.factories import UserFactory
 from project.factories import ProjectFactory
 from notification.factories import AssignmentFactory, NotificationFactory
 from lead.factories import LeadFactory
-from entry.factories import EntryCommentFactory, EntryFactory
+from entry.factories import EntryFactory
+from quality_assurance.factories import EntryReviewCommentFactory
 from analysis_framework.factories import AnalysisFrameworkFactory
 
 
@@ -230,7 +230,11 @@ def test_assignments_query(self):
               totalCount
             }
           }
-'''
+        '''
+
+        # XXX: To avoid using content type cache from pre-tests
+        ContentType.objects.clear_cache()
+
         project = ProjectFactory.create()
         user = UserFactory.create()
         another = UserFactory.create()
@@ -240,17 +244,16 @@ def test_assignments_query(self):
             analysis_framework=af,
             lead=lead
         )
-        entry_comment = EntryCommentFactory.create(
+        entry_comment = EntryReviewCommentFactory.create(
             entry=entry,
             created_by=user
         )
-        Assignment.objects.all().delete()
+
         AssignmentFactory.create_batch(
             3,
             project=project,
-            object_id=lead.id,
-            content_type=ContentType.objects.get_for_model(Lead),
-            created_for=user
+            content_object=lead,
+            created_for=user,
         )
 
         def _query_check(**kwargs):
@@ -271,8 +274,7 @@ def _query_check(**kwargs):
         AssignmentFactory.create_batch(
             3,
             project=project,
-            object_id=entry_comment.id,
-            content_type=ContentType.objects.get_for_model(EntryReviewComment),
+            content_object=entry_comment,
             created_for=user
         )
         content = _query_check()
@@ -280,19 +282,19 @@ def _query_check(**kwargs):
 
     def test_assignments_with_filter_query(self):
         query = '''
-            query MyQuery (
-            $isDone : Boolean,
-        ){
-         assignments(
-            isDone: $isDone,
-        ){
-          totalCount
-          results{
-            id
-        }
-        }
+        query MyQuery($isDone: Boolean) {
+          assignments(isDone: $isDone) {
+            totalCount
+            results {
+              id
+            }
+          }
         }
         '''
+
+        # XXX: To avoid using content type cache from pre-tests
+        ContentType.objects.clear_cache()
+
         project = ProjectFactory.create()
         user = UserFactory.create()
         lead = LeadFactory.create()
@@ -301,24 +303,22 @@ def test_assignments_with_filter_query(self):
             analysis_framework=af,
             lead=lead
         )
-        entry_comment = EntryCommentFactory.create(
+        entry_comment = EntryReviewCommentFactory.create(
             entry=entry,
             created_by=user
         )
-        Assignment.objects.all().delete()
+
         AssignmentFactory.create_batch(
             3,
             project=project,
-            object_id=lead.id,
-            content_type=ContentType.objects.get_for_model(Lead),
+            content_object=lead,
             created_for=user,
             is_done=False
         )
         AssignmentFactory.create_batch(
             5,
             project=project,
-            object_id=entry_comment.id,
-            content_type=ContentType.objects.get_for_model(EntryReviewComment),
+            content_object=entry_comment,
             created_for=user,
             is_done=True
         )