Skip to content

Commit

Permalink
Format tests
Browse files Browse the repository at this point in the history
Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
  • Loading branch information
susodapop committed Jul 10, 2023
1 parent f606e59 commit 1c3d080
Showing 1 changed file with 30 additions and 17 deletions.
47 changes: 30 additions & 17 deletions tests/e2e/sqlalchemy/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,24 +16,28 @@

USER_AGENT_TOKEN = "PySQL e2e Tests"


def sqlalchemy_1_3():
import sqlalchemy

return sqlalchemy.__version__.startswith("1.3")


def version_agnostic_select(object_to_select, *args, **kwargs):
"""
SQLAlchemy==1.3.x requires arguments to select() to be a Python list
https://docs.sqlalchemy.org/en/20/changelog/migration_14.html#orm-query-is-internally-unified-with-select-update-delete-2-0-style-execution-available
"""

if sqlalchemy_1_3():
return select([object_to_select], *args, **kwargs)
else:
return select(object_to_select, *args, **kwargs)



def version_agnostic_connect_arguments(catalog=None, schema=None) -> Tuple[str, dict]:

HOST = os.environ.get("host")
HTTP_PATH = os.environ.get("http_path")
ACCESS_TOKEN = os.environ.get("access_token")
Expand All @@ -44,29 +48,34 @@ def version_agnostic_connect_arguments(catalog=None, schema=None) -> Tuple[str,

if sqlalchemy_1_3():
conn_string = f"databricks://token:{ACCESS_TOKEN}@{HOST}"
connect_args = {**ua_connect_args,
connect_args = {
**ua_connect_args,
"http_path": HTTP_PATH,
"server_hostname": HOST,
"catalog": CATALOG,
"schema": SCHEMA
"schema": SCHEMA,
}

return conn_string, connect_args
else:
return f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}", ua_connect_args


return (
f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}",
ua_connect_args,
)


@pytest.fixture
def db_engine() -> Engine:
conn_string, connect_args = version_agnostic_connect_arguments()
return create_engine(conn_string, connect_args=connect_args)


@pytest.fixture
def samples_engine() -> Engine:

conn_string, connect_args = version_agnostic_connect_arguments(catalog="samples", schema="nyctaxi")
conn_string, connect_args = version_agnostic_connect_arguments(
catalog="samples", schema="nyctaxi"
)
return create_engine(conn_string, connect_args=connect_args)


Expand Down Expand Up @@ -130,7 +139,7 @@ def test_pandas_upload(db_engine, metadata_obj):
db_engine.execute("DROP TABLE mock_data")


def test_create_table_not_null(db_engine, metadata_obj:MetaData):
def test_create_table_not_null(db_engine, metadata_obj: MetaData):

table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))

Expand Down Expand Up @@ -301,29 +310,33 @@ def test_dialect_type_mappings(base, db_engine, metadata_obj: MetaData):

metadata_obj.drop_all()


def test_inspector_smoke_test(samples_engine: Engine):
"""It does not appear that 3L namespace is supported here
"""
"""It does not appear that 3L namespace is supported here"""

from sqlalchemy.engine.reflection import Inspector
schema, table = "nyctaxi", "trips"


schema, table = "nyctaxi", "trips"

try:
inspector = Inspector.from_engine(samples_engine)
except Exception as e:
assert False, f"Could not build inspector: {e}"

# Expect six columns
columns = inspector.get_columns(table, schema=schema)

# Expect zero views, but the method should return
views = inspector.get_view_names(schema=schema)

assert len(columns) == 6, "Dialect did not find the expected number of columns in samples.nyctaxi.trips"
assert (
len(columns) == 6
), "Dialect did not find the expected number of columns in samples.nyctaxi.trips"
assert len(views) == 0, "Views could not be fetched"


def test_get_table_names_smoke_test(samples_engine: Engine):

with samples_engine.connect() as conn:
_names = samples_engine.table_names(schema="nyctaxi", connection=conn)
_names is not None, "get_table_names did not succeed"

0 comments on commit 1c3d080

Please sign in to comment.