Test for custom column types

Author: tomchristie

tests/test_databases.py

@@ -13,6 +13,18 @@
 DATABASE_URLS = [url.strip() for url in os.environ["TEST_DATABASE_URLS"].split(",")]
 
 
+class MyEpochType(sqlalchemy.types.TypeDecorator):
+    impl = sqlalchemy.Integer
+
+    epoch = datetime.date(1970, 1, 1)
+
+    def process_bind_param(self, value, dialect):
+        return (value - self.epoch).days
+
+    def process_result_value(self, value, dialect):
+        return self.epoch + datetime.timedelta(days=value)
+
+
 metadata = sqlalchemy.MetaData()
 
 notes = sqlalchemy.Table(
@@ -23,6 +35,7 @@
     sqlalchemy.Column("completed", sqlalchemy.Boolean),
 )
 
+# Used to test DateTime
 articles = sqlalchemy.Table(
     "articles",
     metadata,
@@ -31,13 +44,23 @@
     sqlalchemy.Column("published", sqlalchemy.DateTime),
 )
 
+# Used to test JSON
 session = sqlalchemy.Table(
     "session",
     metadata,
     sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True),
     sqlalchemy.Column("data", sqlalchemy.JSON),
 )
 
+# Used to test custom column types
+custom_date = sqlalchemy.Table(
+    "custom_date",
+    metadata,
+    sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True),
+    sqlalchemy.Column("title", sqlalchemy.String(length=100)),
+    sqlalchemy.Column("published", MyEpochType),
+)
+
 
 @pytest.fixture(autouse=True, scope="module")
 def create_test_database():
@@ -226,7 +249,7 @@ async def test_transaction_rollback_low_level(database_url):
 @async_adapter
 async def test_datetime_field(database_url):
     """
-    Test DataTime fields, to ensure records are coerced to proper Python types.
+    Test DataTime columns, to ensure records are coerced to/from proper Python types.
     """
 
     async with Database(database_url) as database:
@@ -250,7 +273,7 @@ async def test_datetime_field(database_url):
 @async_adapter
 async def test_json_field(database_url):
     """
-    Test JSON fields, to ensure correct cross-database support.
+    Test JSON columns, to ensure correct cross-database support.
     """
 
     async with Database(database_url) as database:
@@ -265,3 +288,27 @@ async def test_json_field(database_url):
             results = await database.fetch_all(query=query)
             assert len(results) == 1
             assert results[0]["data"] == {"text": "hello", "boolean": True, "int": 1}
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_custom_field(database_url):
+    """
+    Test custom column types.
+    """
+
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            today = datetime.date.today()
+
+            # execute()
+            query = custom_date.insert()
+            values = {"title": "Hello, world", "published": today}
+            await database.execute(query, values)
+
+            # fetch_all()
+            query = custom_date.select()
+            results = await database.fetch_all(query=query)
+            assert len(results) == 1
+            assert results[0]["title"] == "Hello, world"
+            assert results[0]["published"] == today