Updated tests.

xnuinside · Jan 31, 2024 · b2cfa80 · b2cfa80
1 parent 386e00b
commit b2cfa80
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/tests/functional/generator/test_sqlalchemy_core.py b/tests/functional/generator/test_sqlalchemy_core.py
@@ -118,7 +118,7 @@ class Products(Base):
 
     __tablename__ = 'products'
 
-    id = sa.Column(sa.Integer(), nullable=False)
+    id = sa.Column(sa.Integer(), primary_key=True)
     merchant_id = sa.Column(sa.Integer(), sa.ForeignKey('merchants.id'), nullable=False)
 """
 

diff --git a/tests/functional/generator/test_sqlmodel.py b/tests/functional/generator/test_sqlmodel.py
@@ -3,6 +3,7 @@
 
 def test_with_enums():
     expected = """import datetime
+import decimal
 from typing import Optional
 from sqlmodel import Field, SQLModel
 
@@ -56,6 +57,7 @@ class Material(SQLModel, table=True):
 
 def test_foreign_keys():
     expected = """import datetime
+import decimal
 from typing import Optional
 from sqlmodel import Field, SQLModel
 
@@ -133,6 +135,7 @@ def test_foreign_keys_defined_inline():
     different, yet valid input DDL.
     """
     expected = """import datetime
+import decimal
 from typing import Optional
 from sqlmodel import Field, SQLModel
 
@@ -206,6 +209,7 @@ def test_multi_col_pk_and_fk():
     FOREIGN keys.
     """
     expected = """import datetime
+import decimal
 from typing import Optional
 from sqlmodel import Field, SQLModel
 
@@ -260,6 +264,7 @@ class LinkedTo(SQLModel, table=True):
 
 def test_upper_name_produces_the_same_result():
     expected = """import datetime
+import decimal
 from typing import Optional
 from sqlmodel import Field, SQLModel
 
@@ -313,6 +318,7 @@ class Material(SQLModel, table=True):
 
 def test_foreign_keys_in_different_schema():
     expected = """import datetime
+import decimal
 from typing import Optional
 from sqlmodel import Field, SQLModel