Fix test cases for batching in sqlalchemy1.4

Author: colelin26

graphene_sqlalchemy/tests/test_batching.py

@@ -1,3 +1,4 @@
+import ast
 import contextlib
 import logging
 
@@ -9,9 +10,9 @@
 from ..fields import (BatchSQLAlchemyConnectionField,
                       default_connection_field_factory)
 from ..types import ORMField, SQLAlchemyObjectType
-from .models import Article, HairKind, Pet, Reporter
 from ..utils import is_sqlalchemy_version_less_than
-from .utils import to_std_dicts
+from .models import Article, HairKind, Pet, Reporter
+from .utils import remove_cache_miss_stat, to_std_dicts
 
 
 class MockLoggingHandler(logging.Handler):
@@ -127,26 +128,12 @@ async def test_many_to_one(session_factory):
         assert len(sql_statements) == 1
         return
 
-    assert messages == [
-      'BEGIN (implicit)',
-
-      'SELECT articles.id AS articles_id, '
-      'articles.headline AS articles_headline, '
-      'articles.pub_date AS articles_pub_date, '
-      'articles.reporter_id AS articles_reporter_id \n'
-      'FROM articles',
-      '()',
-
-      'SELECT reporters.id AS reporters_id, '
-      '(SELECT CAST(count(reporters.id) AS INTEGER) AS anon_2 \nFROM reporters) AS anon_1, '
-      'reporters.first_name AS reporters_first_name, '
-      'reporters.last_name AS reporters_last_name, '
-      'reporters.email AS reporters_email, '
-      'reporters.favorite_pet_kind AS reporters_favorite_pet_kind \n'
-      'FROM reporters \n'
-      'WHERE reporters.id IN (?, ?)',
-      '(1, 2)',
-    ]
+    if not is_sqlalchemy_version_less_than('1.4'):
+        messages[2] = remove_cache_miss_stat(messages[2])
+        messages[4] = remove_cache_miss_stat(messages[4])
+
+    assert ast.literal_eval(messages[2]) == ()
+    assert sorted(ast.literal_eval(messages[4])) == [1, 2]
 
     assert not result.errors
     result = to_std_dicts(result.data)
@@ -219,26 +206,12 @@ async def test_one_to_one(session_factory):
         assert len(sql_statements) == 1
         return
 
-    assert messages == [
-      'BEGIN (implicit)',
-
-      'SELECT (SELECT CAST(count(reporters.id) AS INTEGER) AS anon_2 \nFROM reporters) AS anon_1, '
-      'reporters.id AS reporters_id, '
-      'reporters.first_name AS reporters_first_name, '
-      'reporters.last_name AS reporters_last_name, '
-      'reporters.email AS reporters_email, '
-      'reporters.favorite_pet_kind AS reporters_favorite_pet_kind \n'
-      'FROM reporters',
-      '()',
-
-      'SELECT articles.reporter_id AS articles_reporter_id, '
-      'articles.id AS articles_id, '
-      'articles.headline AS articles_headline, '
-      'articles.pub_date AS articles_pub_date \n'
-      'FROM articles \n'
-      'WHERE articles.reporter_id IN (?, ?)',
-      '(1, 2)'
-    ]
+    if not is_sqlalchemy_version_less_than('1.4'):
+        messages[2] = remove_cache_miss_stat(messages[2])
+        messages[4] = remove_cache_miss_stat(messages[4])
+
+    assert ast.literal_eval(messages[2]) == ()
+    assert sorted(ast.literal_eval(messages[4])) == [1, 2]
 
     assert not result.errors
     result = to_std_dicts(result.data)
@@ -323,26 +296,12 @@ async def test_one_to_many(session_factory):
         assert len(sql_statements) == 1
         return
 
-    assert messages == [
-      'BEGIN (implicit)',
-
-      'SELECT (SELECT CAST(count(reporters.id) AS INTEGER) AS anon_2 \nFROM reporters) AS anon_1, '
-      'reporters.id AS reporters_id, '
-      'reporters.first_name AS reporters_first_name, '
-      'reporters.last_name AS reporters_last_name, '
-      'reporters.email AS reporters_email, '
-      'reporters.favorite_pet_kind AS reporters_favorite_pet_kind \n'
-      'FROM reporters',
-      '()',
-
-      'SELECT articles.reporter_id AS articles_reporter_id, '
-      'articles.id AS articles_id, '
-      'articles.headline AS articles_headline, '
-      'articles.pub_date AS articles_pub_date \n'
-      'FROM articles \n'
-      'WHERE articles.reporter_id IN (?, ?)',
-      '(1, 2)'
-    ]
+    if not is_sqlalchemy_version_less_than('1.4'):
+        messages[2] = remove_cache_miss_stat(messages[2])
+        messages[4] = remove_cache_miss_stat(messages[4])
+
+    assert ast.literal_eval(messages[2]) == ()
+    assert sorted(ast.literal_eval(messages[4])) == [1, 2]
 
     assert not result.errors
     result = to_std_dicts(result.data)
@@ -451,31 +410,12 @@ async def test_many_to_many(session_factory):
         assert len(sql_statements) == 1
         return
 
-    assert messages == [
-      'BEGIN (implicit)',
-
-      'SELECT (SELECT CAST(count(reporters.id) AS INTEGER) AS anon_2 \nFROM reporters) AS anon_1, '
-      'reporters.id AS reporters_id, '
-      'reporters.first_name AS reporters_first_name, '
-      'reporters.last_name AS reporters_last_name, '
-      'reporters.email AS reporters_email, '
-      'reporters.favorite_pet_kind AS reporters_favorite_pet_kind \n'
-      'FROM reporters',
-      '()',
-
-      'SELECT reporters_1.id AS reporters_1_id, '
-      'pets.id AS pets_id, '
-      'pets.name AS pets_name, '
-      'pets.pet_kind AS pets_pet_kind, '
-      'pets.hair_kind AS pets_hair_kind, '
-      'pets.reporter_id AS pets_reporter_id \n'
-      'FROM reporters AS reporters_1 '
-      'JOIN association AS association_1 ON reporters_1.id = association_1.reporter_id '
-      'JOIN pets ON pets.id = association_1.pet_id \n'
-      'WHERE reporters_1.id IN (?, ?) '
-      'ORDER BY pets.id',
-      '(1, 2)'
-    ]
+    if not is_sqlalchemy_version_less_than('1.4'):
+        messages[2] = remove_cache_miss_stat(messages[2])
+        messages[4] = remove_cache_miss_stat(messages[4])
+
+    assert ast.literal_eval(messages[2]) == ()
+    assert sorted(ast.literal_eval(messages[4])) == [1, 2]
 
     assert not result.errors
     result = to_std_dicts(result.data)

graphene_sqlalchemy/tests/utils.py

@@ -1,3 +1,6 @@
+import re
+
+
 def to_std_dicts(value):
     """Convert nested ordered dicts to normal dicts for better comparison."""
     if isinstance(value, dict):
@@ -6,3 +9,9 @@ def to_std_dicts(value):
         return [to_std_dicts(v) for v in value]
     else:
         return value
+
+
+def remove_cache_miss_stat(message):
+    """Remove the stat from the echoed query message when the cache is missed for sqlalchemy version >= 1.4"""
+    # https://github.com/sqlalchemy/sqlalchemy/blob/990eb3d8813369d3b8a7776ae85fb33627443d30/lib/sqlalchemy/engine/default.py#L1177
+    return re.sub(r"\[generated in \d+.?\d*s\]\s", "", message)