Fix test failures with older Python versions

Author: danielfrankcom

aurora_dsql_django/schema.py

@@ -19,6 +19,7 @@
 
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.backends.postgresql import schema
+from django.db.models import CheckConstraint
 
 
 class DatabaseSchemaEditor(schema.DatabaseSchemaEditor):
@@ -64,13 +65,13 @@ def __enter__(self):
 
     def add_index(self, model, index, concurrently=False):
         if index.contains_expressions and not self.connection.features.supports_expression_indexes:
-            return None
-        return super().add_index(model, index, concurrently)
+            return
+        super().add_index(model, index, concurrently)
 
     def remove_index(self, model, index, concurrently=False):
         if index.contains_expressions and not self.connection.features.supports_expression_indexes:
-            return None
-        return super().remove_index(model, index, concurrently)
+            return
+        super().remove_index(model, index, concurrently)
 
     def _check_sql(self, name, check):
         # There is no feature check in the upstream implementation when creating
@@ -79,6 +80,18 @@ def _check_sql(self, name, check):
             return None
         return super()._check_sql(name, check)
 
+    def add_constraint(self, model, constraint):
+        # Older versions of Django don't reference supports_table_check_constraints, so we add this as a backup.
+        if isinstance(constraint, CheckConstraint) and not self.connection.features.supports_table_check_constraints:
+            return
+        super().add_constraint(model, constraint)
+
+    def remove_constraint(self, model, constraint):
+        # Older versions of Django don't reference supports_table_check_constraints, so we add this as a backup.
+        if isinstance(constraint, CheckConstraint) and not self.connection.features.supports_table_check_constraints:
+            return
+        super().remove_constraint(model, constraint)
+
     def _index_columns(self, table, columns, col_suffixes, opclasses):
         # Aurora DSQL doesn't support PostgreSQL opclasses.
         return BaseDatabaseSchemaEditor._index_columns(

aurora_dsql_django/tests/unit/test_schema.py

@@ -1,7 +1,12 @@
 import unittest
 from unittest.mock import patch, MagicMock
-from aurora_dsql_django.schema import DatabaseSchemaEditor
+
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.models.constraints import UniqueConstraint
+from django.db.models.query_utils import Q
+
+from aurora_dsql_django.schema import DatabaseSchemaEditor
+from aurora_dsql_django.tests.utils import create_check_constraint
 
 
 class TestDatabaseSchemaEditor(unittest.TestCase):
@@ -110,6 +115,68 @@ def test_check_sql_feature_enabled(self, mock_super_check_sql):
         mock_super_check_sql.assert_called_once_with("test_check", "age >= 0")
         self.assertEqual(result, "CHECK (age >= 0)")
 
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor.add_constraint')
+    def test_add_constraint_check_disabled(self, mock_super_add_constraint):
+        model = MagicMock()
+        constraint = create_check_constraint(Q(age__gte=0), 'age_check')
+        self.connection.features.supports_table_check_constraints = False
+        
+        result = self.schema_editor.add_constraint(model, constraint)
+        
+        self.assertIsNone(result)
+        mock_super_add_constraint.assert_not_called()
+
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor.add_constraint')
+    def test_add_constraint_check_enabled(self, mock_super_add_constraint):
+        model = MagicMock()
+        constraint = create_check_constraint(Q(age__gte=0), 'age_check')
+        self.connection.features.supports_table_check_constraints = True
+        
+        self.schema_editor.add_constraint(model, constraint)
+        
+        mock_super_add_constraint.assert_called_once_with(model, constraint)
+
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor.add_constraint')
+    def test_add_constraint_non_check(self, mock_super_add_constraint):
+        model = MagicMock()
+        constraint = UniqueConstraint(fields=['name'], name='unique_name')
+        self.connection.features.supports_table_check_constraints = False
+        
+        self.schema_editor.add_constraint(model, constraint)
+        
+        mock_super_add_constraint.assert_called_once_with(model, constraint)
+
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor.remove_constraint')
+    def test_remove_constraint_check_disabled(self, mock_super_remove_constraint):
+        model = MagicMock()
+        constraint = create_check_constraint(Q(age__gte=0), 'age_check')
+        self.connection.features.supports_table_check_constraints = False
+        
+        result = self.schema_editor.remove_constraint(model, constraint)
+        
+        self.assertIsNone(result)
+        mock_super_remove_constraint.assert_not_called()
+
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor.remove_constraint')
+    def test_remove_constraint_check_enabled(self, mock_super_remove_constraint):
+        model = MagicMock()
+        constraint = create_check_constraint(Q(age__gte=0), 'age_check')
+        self.connection.features.supports_table_check_constraints = True
+        
+        self.schema_editor.remove_constraint(model, constraint)
+        
+        mock_super_remove_constraint.assert_called_once_with(model, constraint)
+
+    @patch('aurora_dsql_django.schema.schema.DatabaseSchemaEditor.remove_constraint')
+    def test_remove_constraint_non_check(self, mock_super_remove_constraint):
+        model = MagicMock()
+        constraint = UniqueConstraint(fields=['name'], name='unique_name')
+        self.connection.features.supports_table_check_constraints = False
+        
+        self.schema_editor.remove_constraint(model, constraint)
+        
+        mock_super_remove_constraint.assert_called_once_with(model, constraint)
+
 
 if __name__ == '__main__':
     unittest.main()

aurora_dsql_django/tests/unit/test_wrapper.py

@@ -7,6 +7,7 @@
 from aurora_dsql_django.base import DatabaseWrapper
 from aurora_dsql_django.features import DatabaseFeatures
 from aurora_dsql_django.schema import DatabaseSchemaEditor
+from aurora_dsql_django.tests.utils import create_check_constraint
 
 if not settings.configured:
     settings.configure(
@@ -87,7 +88,7 @@ class CheckConstraintModel(models.Model):
             class Meta:
                 app_label = 'test_app'
                 constraints = [
-                    CheckConstraint(condition=Q(age__gte=0), name='age_gte_0')
+                    create_check_constraint(Q(age__gte=0), 'age_gte_0')
                 ]
 
         def operation():
@@ -109,7 +110,7 @@ class AddCheckConstraintModel(models.Model):
             class Meta:
                 app_label = 'test_app'
 
-        constraint = CheckConstraint(condition=Q(age__gte=0), name='age_gte_0')
+        constraint = create_check_constraint(Q(age__gte=0), 'age_gte_0')
 
         def operation():
             self.schema_editor.add_constraint(AddCheckConstraintModel, constraint)
@@ -129,7 +130,7 @@ class RemoveCheckConstraintModel(models.Model):
             class Meta:
                 app_label = 'test_app'
 
-        constraint = CheckConstraint(condition=Q(age__gte=0), name='age_gte_0')
+        constraint = create_check_constraint(Q(age__gte=0), 'age_gte_0')
 
         def operation():
             self.schema_editor.remove_constraint(RemoveCheckConstraintModel, constraint)

aurora_dsql_django/tests/utils.py

@@ -0,0 +1,13 @@
+"""Shared test utilities for Aurora DSQL Django tests."""
+
+import django
+from django.db.models import CheckConstraint
+
+
+# noinspection PyArgumentList
+def create_check_constraint(condition, name):
+    """Create CheckConstraint with Django version compatibility"""
+    if django.VERSION >= (5, 0):
+        return CheckConstraint(condition=condition, name=name)
+    else:
+        return CheckConstraint(check=condition, name=name)