Skip to content

Commit dae15e3

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
formatting (black)
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent ed7cf91 commit dae15e3

6 files changed

Lines changed: 133 additions & 73 deletions

File tree

‎examples/experimental/sea_connector_test.py‎

Lines changed: 31 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -22,90 +22,99 @@
2222
"test_sea_metadata",
2323
]
2424

25+
2526
def load_test_function(module_name: str) -> Callable:
2627
"""Load a test function from a module."""
2728
module_path = os.path.join(
28-
os.path.dirname(os.path.abspath(__file__)),
29-
"tests",
30-
f"{module_name}.py"
29+
os.path.dirname(os.path.abspath(__file__)), "tests", f"{module_name}.py"
3130
)
32-
31+
3332
spec = importlib.util.spec_from_file_location(module_name, module_path)
3433
module = importlib.util.module_from_spec(spec)
3534
spec.loader.exec_module(module)
36-
35+
3736
# Get the main test function (assuming it starts with "test_")
3837
for name in dir(module):
3938
if name.startswith("test_") and callable(getattr(module, name)):
4039
# For sync and async query modules, we want the main function that runs both tests
4140
if name == f"test_sea_{module_name.replace('test_sea_', '')}_exec":
4241
return getattr(module, name)
43-
42+
4443
# Fallback to the first test function found
4544
for name in dir(module):
4645
if name.startswith("test_") and callable(getattr(module, name)):
4746
return getattr(module, name)
48-
47+
4948
raise ValueError(f"No test function found in module {module_name}")
5049

50+
5151
def run_tests() -> List[Tuple[str, bool]]:
5252
"""Run all tests and return results."""
5353
results = []
54-
54+
5555
for module_name in TEST_MODULES:
5656
try:
5757
test_func = load_test_function(module_name)
5858
logger.info(f"\n{'=' * 50}")
5959
logger.info(f"Running test: {module_name}")
6060
logger.info(f"{'-' * 50}")
61-
61+
6262
success = test_func()
6363
results.append((module_name, success))
64-
64+
6565
status = "✅ PASSED" if success else "❌ FAILED"
6666
logger.info(f"Test {module_name}: {status}")
67-
67+
6868
except Exception as e:
6969
logger.error(f"Error loading or running test {module_name}: {str(e)}")
7070
import traceback
71+
7172
logger.error(traceback.format_exc())
7273
results.append((module_name, False))
73-
74+
7475
return results
7576

77+
7678
def print_summary(results: List[Tuple[str, bool]]) -> None:
7779
"""Print a summary of test results."""
7880
logger.info(f"\n{'=' * 50}")
7981
logger.info("TEST SUMMARY")
8082
logger.info(f"{'-' * 50}")
81-
83+
8284
passed = sum(1 for _, success in results if success)
8385
total = len(results)
84-
86+
8587
for module_name, success in results:
8688
status = "✅ PASSED" if success else "❌ FAILED"
8789
logger.info(f"{status} - {module_name}")
88-
90+
8991
logger.info(f"{'-' * 50}")
9092
logger.info(f"Total: {total} | Passed: {passed} | Failed: {total - passed}")
9193
logger.info(f"{'=' * 50}")
9294

95+
9396
if __name__ == "__main__":
9497
# Check if required environment variables are set
95-
required_vars = ["DATABRICKS_SERVER_HOSTNAME", "DATABRICKS_HTTP_PATH", "DATABRICKS_TOKEN"]
98+
required_vars = [
99+
"DATABRICKS_SERVER_HOSTNAME",
100+
"DATABRICKS_HTTP_PATH",
101+
"DATABRICKS_TOKEN",
102+
]
96103
missing_vars = [var for var in required_vars if not os.environ.get(var)]
97-
104+
98105
if missing_vars:
99-
logger.error(f"Missing required environment variables: {', '.join(missing_vars)}")
106+
logger.error(
107+
f"Missing required environment variables: {', '.join(missing_vars)}"
108+
)
100109
logger.error("Please set these variables before running the tests.")
101110
sys.exit(1)
102-
111+
103112
# Run all tests
104113
results = run_tests()
105-
114+
106115
# Print summary
107116
print_summary(results)
108-
117+
109118
# Exit with appropriate status code
110119
all_passed = all(success for _, success in results)
111-
sys.exit(0 if all_passed else 1)
120+
sys.exit(0 if all_passed else 1)

‎examples/experimental/tests/__init__.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +0,0 @@
1-
# This file makes the tests directory a Python package

‎examples/experimental/tests/test_sea_async_query.py‎

Lines changed: 49 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,9 @@ def test_sea_async_query_with_cloud_fetch():
3333

3434
try:
3535
# Create connection with cloud fetch enabled
36-
logger.info("Creating connection for asynchronous query execution with cloud fetch enabled")
36+
logger.info(
37+
"Creating connection for asynchronous query execution with cloud fetch enabled"
38+
)
3739
connection = Connection(
3840
server_hostname=server_hostname,
3941
http_path=http_path,
@@ -51,30 +53,39 @@ def test_sea_async_query_with_cloud_fetch():
5153

5254
# Execute a simple query asynchronously
5355
cursor = connection.cursor()
54-
logger.info("Executing asynchronous query with cloud fetch: SELECT 1 as test_value")
56+
logger.info(
57+
"Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
58+
)
5559
cursor.execute_async("SELECT 1 as test_value")
56-
logger.info("Asynchronous query submitted successfully with cloud fetch enabled")
57-
60+
logger.info(
61+
"Asynchronous query submitted successfully with cloud fetch enabled"
62+
)
63+
5864
# Check query state
5965
logger.info("Checking query state...")
6066
while cursor.is_query_pending():
6167
logger.info("Query is still pending, waiting...")
6268
time.sleep(1)
63-
69+
6470
logger.info("Query is no longer pending, getting results...")
6571
cursor.get_async_execution_result()
66-
logger.info("Successfully retrieved asynchronous query results with cloud fetch enabled")
67-
72+
logger.info(
73+
"Successfully retrieved asynchronous query results with cloud fetch enabled"
74+
)
75+
6876
# Close resources
6977
cursor.close()
7078
connection.close()
7179
logger.info("Successfully closed SEA session")
72-
80+
7381
return True
7482

7583
except Exception as e:
76-
logger.error(f"Error during SEA asynchronous query execution test with cloud fetch: {str(e)}")
84+
logger.error(
85+
f"Error during SEA asynchronous query execution test with cloud fetch: {str(e)}"
86+
)
7787
import traceback
88+
7889
logger.error(traceback.format_exc())
7990
return False
8091

@@ -100,7 +111,9 @@ def test_sea_async_query_without_cloud_fetch():
100111

101112
try:
102113
# Create connection with cloud fetch disabled
103-
logger.info("Creating connection for asynchronous query execution with cloud fetch disabled")
114+
logger.info(
115+
"Creating connection for asynchronous query execution with cloud fetch disabled"
116+
)
104117
connection = Connection(
105118
server_hostname=server_hostname,
106119
http_path=http_path,
@@ -119,30 +132,39 @@ def test_sea_async_query_without_cloud_fetch():
119132

120133
# Execute a simple query asynchronously
121134
cursor = connection.cursor()
122-
logger.info("Executing asynchronous query without cloud fetch: SELECT 1 as test_value")
135+
logger.info(
136+
"Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
137+
)
123138
cursor.execute_async("SELECT 1 as test_value")
124-
logger.info("Asynchronous query submitted successfully with cloud fetch disabled")
125-
139+
logger.info(
140+
"Asynchronous query submitted successfully with cloud fetch disabled"
141+
)
142+
126143
# Check query state
127144
logger.info("Checking query state...")
128145
while cursor.is_query_pending():
129146
logger.info("Query is still pending, waiting...")
130147
time.sleep(1)
131-
148+
132149
logger.info("Query is no longer pending, getting results...")
133150
cursor.get_async_execution_result()
134-
logger.info("Successfully retrieved asynchronous query results with cloud fetch disabled")
135-
151+
logger.info(
152+
"Successfully retrieved asynchronous query results with cloud fetch disabled"
153+
)
154+
136155
# Close resources
137156
cursor.close()
138157
connection.close()
139158
logger.info("Successfully closed SEA session")
140-
159+
141160
return True
142161

143162
except Exception as e:
144-
logger.error(f"Error during SEA asynchronous query execution test without cloud fetch: {str(e)}")
163+
logger.error(
164+
f"Error during SEA asynchronous query execution test without cloud fetch: {str(e)}"
165+
)
145166
import traceback
167+
146168
logger.error(traceback.format_exc())
147169
return False
148170

@@ -152,14 +174,18 @@ def test_sea_async_query_exec():
152174
Run both asynchronous query tests and return overall success.
153175
"""
154176
with_cloud_fetch_success = test_sea_async_query_with_cloud_fetch()
155-
logger.info(f"Asynchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}")
156-
177+
logger.info(
178+
f"Asynchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}"
179+
)
180+
157181
without_cloud_fetch_success = test_sea_async_query_without_cloud_fetch()
158-
logger.info(f"Asynchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}")
159-
182+
logger.info(
183+
f"Asynchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}"
184+
)
185+
160186
return with_cloud_fetch_success and without_cloud_fetch_success
161187

162188

163189
if __name__ == "__main__":
164190
success = test_sea_async_query_exec()
165-
sys.exit(0 if success else 1)
191+
sys.exit(0 if success else 1)

‎examples/experimental/tests/test_sea_metadata.py‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,11 @@ def test_sea_metadata():
2828
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
2929
)
3030
return False
31-
31+
3232
if not catalog:
33-
logger.error("DATABRICKS_CATALOG environment variable is required for metadata tests.")
33+
logger.error(
34+
"DATABRICKS_CATALOG environment variable is required for metadata tests."
35+
)
3436
return False
3537

3638
try:
@@ -55,37 +57,42 @@ def test_sea_metadata():
5557
logger.info("Fetching catalogs...")
5658
cursor.catalogs()
5759
logger.info("Successfully fetched catalogs")
58-
60+
5961
# Test schemas
6062
logger.info(f"Fetching schemas for catalog '{catalog}'...")
6163
cursor.schemas(catalog_name=catalog)
6264
logger.info("Successfully fetched schemas")
63-
65+
6466
# Test tables
6567
logger.info(f"Fetching tables for catalog '{catalog}', schema 'default'...")
6668
cursor.tables(catalog_name=catalog, schema_name="default")
6769
logger.info("Successfully fetched tables")
68-
70+
6971
# Test columns for a specific table
7072
# Using a common table that should exist in most environments
71-
logger.info(f"Fetching columns for catalog '{catalog}', schema 'default', table 'information_schema'...")
72-
cursor.columns(catalog_name=catalog, schema_name="default", table_name="information_schema")
73+
logger.info(
74+
f"Fetching columns for catalog '{catalog}', schema 'default', table 'information_schema'..."
75+
)
76+
cursor.columns(
77+
catalog_name=catalog, schema_name="default", table_name="information_schema"
78+
)
7379
logger.info("Successfully fetched columns")
74-
80+
7581
# Close resources
7682
cursor.close()
7783
connection.close()
7884
logger.info("Successfully closed SEA session")
79-
85+
8086
return True
8187

8288
except Exception as e:
8389
logger.error(f"Error during SEA metadata test: {str(e)}")
8490
import traceback
91+
8592
logger.error(traceback.format_exc())
8693
return False
8794

8895

8996
if __name__ == "__main__":
9097
success = test_sea_metadata()
91-
sys.exit(0 if success else 1)
98+
sys.exit(0 if success else 1)

‎examples/experimental/tests/test_sea_session.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,16 +55,17 @@ def test_sea_session():
5555
logger.info("Closing the SEA session...")
5656
connection.close()
5757
logger.info("Successfully closed SEA session")
58-
58+
5959
return True
6060

6161
except Exception as e:
6262
logger.error(f"Error testing SEA session: {str(e)}")
6363
import traceback
64+
6465
logger.error(traceback.format_exc())
6566
return False
6667

6768

6869
if __name__ == "__main__":
6970
success = test_sea_session()
70-
sys.exit(0 if success else 1)
71+
sys.exit(0 if success else 1)

0 commit comments

Comments
 (0)