Black SQLAlchemy tests

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse Whitehouse

tests/e2e/sqlalchemy/test_basic.py

@@ -8,6 +8,7 @@
 
 USER_AGENT_TOKEN = "PySQL e2e Tests"
 
+
 @pytest.fixture
 def db_engine():
 
@@ -19,7 +20,10 @@ def db_engine():
 
     connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}
 
-    engine = create_engine(f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}", connect_args=connect_args)
+    engine = create_engine(
+        f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}",
+        connect_args=connect_args,
+    )
     return engine
 
 
@@ -32,6 +36,7 @@ def base(db_engine):
 def session(db_engine):
     return Session(bind=db_engine)
 
+
 @pytest.fixture()
 def metadata_obj(db_engine):
     return MetaData(bind=db_engine)
@@ -42,6 +47,7 @@ def test_can_connect(db_engine):
     result = db_engine.execute(simple_query).fetchall()
     assert len(result) == 1
 
+
 def test_connect_args(db_engine):
     """Verify that extra connect args passed to sqlalchemy.create_engine are passed to DBAPI
 
@@ -63,7 +69,14 @@ def test_pandas_upload(db_engine, metadata_obj):
     SCHEMA = os.environ.get("schema")
     try:
         df = pd.read_excel("tests/sqlalchemy/demo_data/MOCK_DATA.xlsx")
-        df.to_sql("mock_data", db_engine, schema=SCHEMA, index=False, method="multi", if_exists="replace")
+        df.to_sql(
+            "mock_data",
+            db_engine,
+            schema=SCHEMA,
+            index=False,
+            method="multi",
+            if_exists="replace",
+        )
 
         df_after = pd.read_sql_table("mock_data", db_engine, schema=SCHEMA)
         assert len(df) == len(df_after)
@@ -72,21 +85,24 @@ def test_pandas_upload(db_engine, metadata_obj):
     finally:
         db_engine.execute("DROP TABLE mock_data")
 
+
 def test_create_table_not_null(db_engine, metadata_obj):
 
     table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
 
     SampleTable = Table(
-         table_name,
-         metadata_obj,
-         Column("name", String(255)),
-         Column("episodes", Integer),
-         Column("some_bool", BOOLEAN, nullable=False)
+        table_name,
+        metadata_obj,
+        Column("name", String(255)),
+        Column("episodes", Integer),
+        Column("some_bool", BOOLEAN, nullable=False),
     )
 
     metadata_obj.create_all()
 
-    columns = db_engine.dialect.get_columns(connection=db_engine.connect(), table_name=table_name)
+    columns = db_engine.dialect.get_columns(
+        connection=db_engine.connect(), table_name=table_name
+    )
 
     name_column_description = columns[0]
     some_bool_column_description = columns[2]
@@ -96,6 +112,7 @@ def test_create_table_not_null(db_engine, metadata_obj):
 
     metadata_obj.drop_all()
 
+
 def test_bulk_insert_with_core(db_engine, metadata_obj, session):
 
     import random
@@ -105,39 +122,41 @@ def test_bulk_insert_with_core(db_engine, metadata_obj, session):
     table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
 
     names = ["Bim", "Miki", "Sarah", "Ira"]
-    
+
     SampleTable = Table(
-         table_name,
-         metadata_obj,
-         Column("name", String(255)),
-         Column("number", Integer)
+        table_name, metadata_obj, Column("name", String(255)), Column("number", Integer)
     )
 
-    rows = [{"name": names[i%3], "number": random.choice(range(10000))} for i in range(num_to_insert)]
+    rows = [
+        {"name": names[i % 3], "number": random.choice(range(10000))}
+        for i in range(num_to_insert)
+    ]
 
     metadata_obj.create_all()
     db_engine.execute(insert(SampleTable).values(rows))
 
     rows = db_engine.execute(select(SampleTable)).fetchall()
 
     assert len(rows) == num_to_insert
-    
+
+
 def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData):
-    """
-    """
+    """ """
 
     SampleTable = Table(
-         "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
-         metadata_obj,
-         Column("name", String(255)),
-         Column("episodes", Integer),
-         Column("some_bool", BOOLEAN),
-         Column("dollars", DECIMAL(10,2))
+        "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
+        metadata_obj,
+        Column("name", String(255)),
+        Column("episodes", Integer),
+        Column("some_bool", BOOLEAN),
+        Column("dollars", DECIMAL(10, 2)),
     )
 
     metadata_obj.create_all()
 
-    insert_stmt = insert(SampleTable).values(name="Bim Adewunmi", episodes=6, some_bool=True, dollars=decimal.Decimal(125))
+    insert_stmt = insert(SampleTable).values(
+        name="Bim Adewunmi", episodes=6, some_bool=True, dollars=decimal.Decimal(125)
+    )
 
     with db_engine.connect() as conn:
         conn.execute(insert_stmt)
@@ -156,9 +175,9 @@ def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData):
 # https://docs.sqlalchemy.org/en/14/orm/quickstart.html
 
 
-@skipIf(False, 'Unity catalog must be supported')
+@skipIf(False, "Unity catalog must be supported")
 def test_create_insert_drop_table_orm(base, session: Session):
-    """ORM classes built on the declarative base class must have a primary key. 
+    """ORM classes built on the declarative base class must have a primary key.
     This is restricted to Unity Catalog.
     """
 
@@ -169,7 +188,7 @@ class SampleObject(base):
         name = Column(String(255), primary_key=True)
         episodes = Column(Integer)
         some_bool = Column(BOOLEAN)
-    
+
     base.metadata.create_all()
 
     sample_object_1 = SampleObject(name="Bim Adewunmi", episodes=6, some_bool=True)
@@ -178,7 +197,9 @@ class SampleObject(base):
     session.add(sample_object_2)
     session.commit()
 
-    stmt = select(SampleObject).where(SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"]))
+    stmt = select(SampleObject).where(
+        SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"])
+    )
 
     output = [i for i in session.scalars(stmt)]
     assert len(output) == 2
@@ -187,28 +208,33 @@ class SampleObject(base):
 
 
 def test_dialect_type_mappings(base, db_engine, metadata_obj: MetaData):
-    """Confirms that we get back the same time we declared in a model and inserted using Core
-    """
+    """Confirms that we get back the same time we declared in a model and inserted using Core"""
 
     SampleTable = Table(
-         "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
-         metadata_obj,
-         Column("string_example", String(255)),
-         Column("integer_example", Integer),
-         Column("boolean_example", BOOLEAN),
-         Column("decimal_example", DECIMAL(10,2)),
-         Column("date_example", Date)
+        "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
+        metadata_obj,
+        Column("string_example", String(255)),
+        Column("integer_example", Integer),
+        Column("boolean_example", BOOLEAN),
+        Column("decimal_example", DECIMAL(10, 2)),
+        Column("date_example", Date),
     )
 
     string_example = ""
     integer_example = 100
     boolean_example = True
     decimal_example = decimal.Decimal(125)
-    date_example = datetime.date(2013,1,1)
+    date_example = datetime.date(2013, 1, 1)
 
     metadata_obj.create_all()
 
-    insert_stmt = insert(SampleTable).values(string_example=string_example, integer_example=integer_example, boolean_example=boolean_example, decimal_example=decimal_example, date_example=date_example)
+    insert_stmt = insert(SampleTable).values(
+        string_example=string_example,
+        integer_example=integer_example,
+        boolean_example=boolean_example,
+        decimal_example=decimal_example,
+        date_example=date_example,
+    )
 
     with db_engine.connect() as conn:
         conn.execute(insert_stmt)
@@ -225,4 +251,4 @@ def test_dialect_type_mappings(base, db_engine, metadata_obj: MetaData):
     assert this_row["decimal_example"] == decimal_example
     assert this_row["date_example"] == date_example
 
-    metadata_obj.drop_all()
+    metadata_obj.drop_all()