|
1 | | -""" |
2 | | -Main script to run all SEA connector tests. |
3 | | -
|
4 | | -This script imports and runs all the individual test modules and displays |
5 | | -a summary of test results with visual indicators. |
6 | | -""" |
7 | 1 | importos |
8 | 2 | importsys |
9 | 3 | importlogging |
10 | | -importimportlib.util |
11 | | -fromtypingimportDict,Callable,List,Tuple |
| 4 | +fromdatabricks.sql.clientimportConnection |
12 | 5 |
|
13 | | -# Configure logging |
14 | | -logging.basicConfig(level=logging.INFO) |
| 6 | +logging.basicConfig(level=logging.DEBUG) |
15 | 7 | logger=logging.getLogger(__name__) |
16 | 8 |
|
17 | | -# Define test modules and their main test functions |
18 | | -TEST_MODULES= [ |
19 | | -"test_sea_session", |
20 | | -"test_sea_sync_query", |
21 | | -"test_sea_async_query", |
22 | | -"test_sea_metadata", |
23 | | -] |
24 | | - |
25 | | - |
26 | | -defload_test_function(module_name:str)->Callable: |
27 | | -"""Load a test function from a module.""" |
28 | | -module_path=os.path.join( |
29 | | -os.path.dirname(os.path.abspath(__file__)),"tests",f"{module_name}.py" |
30 | | - ) |
31 | | - |
32 | | -spec=importlib.util.spec_from_file_location(module_name,module_path) |
33 | | -module=importlib.util.module_from_spec(spec) |
34 | | -spec.loader.exec_module(module) |
35 | | - |
36 | | -# Get the main test function (assuming it starts with "test_") |
37 | | -fornameindir(module): |
38 | | -ifname.startswith("test_")andcallable(getattr(module,name)): |
39 | | -# For sync and async query modules, we want the main function that runs both tests |
40 | | -ifname==f"test_sea_{module_name.replace('test_sea_','')}_exec": |
41 | | -returngetattr(module,name) |
42 | | - |
43 | | -# Fallback to the first test function found |
44 | | -fornameindir(module): |
45 | | -ifname.startswith("test_")andcallable(getattr(module,name)): |
46 | | -returngetattr(module,name) |
47 | | - |
48 | | -raiseValueError(f"No test function found in module{module_name}") |
49 | | - |
50 | | - |
51 | | -defrun_tests()->List[Tuple[str,bool]]: |
52 | | -"""Run all tests and return results.""" |
53 | | -results= [] |
54 | | - |
55 | | -formodule_nameinTEST_MODULES: |
56 | | -try: |
57 | | -test_func=load_test_function(module_name) |
58 | | -logger.info(f"\n{'='*50}") |
59 | | -logger.info(f"Running test:{module_name}") |
60 | | -logger.info(f"{'-'*50}") |
61 | | - |
62 | | -success=test_func() |
63 | | -results.append((module_name,success)) |
64 | | - |
65 | | -status="✅ PASSED"ifsuccesselse"❌ FAILED" |
66 | | -logger.info(f"Test{module_name}:{status}") |
67 | | - |
68 | | -exceptExceptionase: |
69 | | -logger.error(f"Error loading or running test{module_name}:{str(e)}") |
70 | | -importtraceback |
71 | | - |
72 | | -logger.error(traceback.format_exc()) |
73 | | -results.append((module_name,False)) |
74 | | - |
75 | | -returnresults |
76 | | - |
77 | | - |
78 | | -defprint_summary(results:List[Tuple[str,bool]])->None: |
79 | | -"""Print a summary of test results.""" |
80 | | -logger.info(f"\n{'='*50}") |
81 | | -logger.info("TEST SUMMARY") |
82 | | -logger.info(f"{'-'*50}") |
83 | | - |
84 | | -passed=sum(1for_,successinresultsifsuccess) |
85 | | -total=len(results) |
86 | | - |
87 | | -formodule_name,successinresults: |
88 | | -status="✅ PASSED"ifsuccesselse"❌ FAILED" |
89 | | -logger.info(f"{status} -{module_name}") |
90 | | - |
91 | | -logger.info(f"{'-'*50}") |
92 | | -logger.info(f"Total:{total} | Passed:{passed} | Failed:{total-passed}") |
93 | | -logger.info(f"{'='*50}") |
94 | | - |
95 | | - |
96 | | -if__name__=="__main__": |
97 | | -# Check if required environment variables are set |
98 | | -required_vars= [ |
99 | | -"DATABRICKS_SERVER_HOSTNAME", |
100 | | -"DATABRICKS_HTTP_PATH", |
101 | | -"DATABRICKS_TOKEN", |
102 | | - ] |
103 | | -missing_vars= [varforvarinrequired_varsifnotos.environ.get(var)] |
104 | | - |
105 | | -ifmissing_vars: |
106 | | -logger.error( |
107 | | -f"Missing required environment variables:{', '.join(missing_vars)}" |
| 9 | +deftest_sea_session(): |
| 10 | +""" |
| 11 | + Test opening and closing a SEA session using the connector. |
| 12 | +
|
| 13 | + This function connects to a Databricks SQL endpoint using the SEA backend, |
| 14 | + opens a session, and then closes it. |
| 15 | +
|
| 16 | + Required environment variables: |
| 17 | + - DATABRICKS_SERVER_HOSTNAME: Databricks server hostname |
| 18 | + - DATABRICKS_HTTP_PATH: HTTP path for the SQL endpoint |
| 19 | + - DATABRICKS_TOKEN: Personal access token for authentication |
| 20 | + """ |
| 21 | + |
| 22 | +server_hostname=os.environ.get("DATABRICKS_SERVER_HOSTNAME") |
| 23 | +http_path=os.environ.get("DATABRICKS_HTTP_PATH") |
| 24 | +access_token=os.environ.get("DATABRICKS_TOKEN") |
| 25 | +catalog=os.environ.get("DATABRICKS_CATALOG") |
| 26 | + |
| 27 | +ifnotall([server_hostname,http_path,access_token]): |
| 28 | +logger.error("Missing required environment variables.") |
| 29 | +logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.") |
| 30 | +sys.exit(1) |
| 31 | + |
| 32 | +logger.info(f"Connecting to{server_hostname}") |
| 33 | +logger.info(f"HTTP Path:{http_path}") |
| 34 | +ifcatalog: |
| 35 | +logger.info(f"Using catalog:{catalog}") |
| 36 | + |
| 37 | +try: |
| 38 | +logger.info("Creating connection with SEA backend...") |
| 39 | +connection=Connection( |
| 40 | +server_hostname=server_hostname, |
| 41 | +http_path=http_path, |
| 42 | +access_token=access_token, |
| 43 | +catalog=catalog, |
| 44 | +schema="default", |
| 45 | +use_sea=True, |
| 46 | +user_agent_entry="SEA-Test-Client"# add custom user agent |
108 | 47 | ) |
109 | | -logger.error("Please set these variables before running the tests.") |
| 48 | + |
| 49 | +logger.info(f"Successfully opened SEA session with ID:{connection.get_session_id_hex()}") |
| 50 | +logger.info(f"backend type:{type(connection.session.backend)}") |
| 51 | + |
| 52 | +# Close the connection |
| 53 | +logger.info("Closing the SEA session...") |
| 54 | +connection.close() |
| 55 | +logger.info("Successfully closed SEA session") |
| 56 | + |
| 57 | +exceptExceptionase: |
| 58 | +logger.error(f"Error testing SEA session:{str(e)}") |
| 59 | +importtraceback |
| 60 | +logger.error(traceback.format_exc()) |
110 | 61 | sys.exit(1) |
| 62 | + |
| 63 | +logger.info("SEA session test completed successfully") |
111 | 64 |
|
112 | | -# Run all tests |
113 | | -results=run_tests() |
114 | | - |
115 | | -# Print summary |
116 | | -print_summary(results) |
117 | | - |
118 | | -# Exit with appropriate status code |
119 | | -all_passed=all(successfor_,successinresults) |
120 | | -sys.exit(0ifall_passedelse1) |
| 65 | +if__name__=="__main__": |
| 66 | +test_sea_session() |