1 files changed, 8 insertions, 2 deletions
diff --git a/mediagoblin/tests/test_submission.py b/mediagoblin/tests/test_submission.py
index c0c3d3cf..f9031d37 100644
--- a/mediagoblin/tests/test_submission.py
+++ b/mediagoblin/tests/test_submission.py
@@ -99,8 +99,14 @@ class TestSubmission:
         return {'upload_files': [('file', filename)]}
 
     def check_comments(self, request, media_id, count):
-        comments = request.db.MediaComment.query.filter_by(media_entry=media_id)
-        assert count == len(list(comments))
+        gmr = request.db.GenericModelReference.query.filter_by(
+            obj_pk=media_id,
+            model_type=request.db.MediaEntry.__tablename__
+        ).first()
+        if gmr is None and count <= 0:
+            return # Yerp it's fine.
+        comments = request.db.Comment.query.filter_by(target_id=gmr.id)
+        assert count == comments.count()
 
     def test_missing_fields(self):
         # Test blank form