|
10 | 10 | import subprocess |
11 | 11 | from typing import List, Tuple |
12 | 12 |
|
13 | | -logging.basicConfig(level=logging.DEBUG) |
| 13 | +# Configure logging |
| 14 | +logging.basicConfig(level=logging.INFO) |
14 | 15 | logger = logging.getLogger(__name__) |
15 | 16 |
|
16 | 17 | TEST_MODULES = [ |
@@ -87,48 +88,29 @@ def print_summary(results: List[Tuple[str, bool]]) -> None: |
87 | 88 | logger.info(f"Total: {total} | Passed: {passed} | Failed: {total - passed}") |
88 | 89 | logger.info(f"{'=' * 50}") |
89 | 90 |
|
90 | | - server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME") |
91 | | - http_path = os.environ.get("DATABRICKS_HTTP_PATH") |
92 | | - access_token = os.environ.get("DATABRICKS_TOKEN") |
93 | | - catalog = os.environ.get("DATABRICKS_CATALOG") |
94 | | - |
95 | | - if not all([server_hostname, http_path, access_token]): |
96 | | - logger.error("Missing required environment variables.") |
97 | | - logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.") |
98 | | - sys.exit(1) |
99 | | - |
100 | | - logger.info(f"Connecting to {server_hostname}") |
101 | | - logger.info(f"HTTP Path: {http_path}") |
102 | | - if catalog: |
103 | | - logger.info(f"Using catalog: {catalog}") |
104 | | - |
105 | | - try: |
106 | | - logger.info("Creating connection with SEA backend...") |
107 | | - connection = Connection( |
108 | | - server_hostname=server_hostname, |
109 | | - http_path=http_path, |
110 | | - access_token=access_token, |
111 | | - catalog=catalog, |
112 | | - schema="default", |
113 | | - use_sea=True, |
114 | | - user_agent_entry="SEA-Test-Client" # add custom user agent |
| 91 | + |
| 92 | +if __name__ == "__main__": |
| 93 | + # Check if required environment variables are set |
| 94 | + required_vars = [ |
| 95 | + "DATABRICKS_SERVER_HOSTNAME", |
| 96 | + "DATABRICKS_HTTP_PATH", |
| 97 | + "DATABRICKS_TOKEN", |
| 98 | + ] |
| 99 | + missing_vars = [var for var in required_vars if not os.environ.get(var)] |
| 100 | + |
| 101 | + if missing_vars: |
| 102 | + logger.error( |
| 103 | + f"Missing required environment variables: {', '.join(missing_vars)}" |
115 | 104 | ) |
116 | | - |
117 | | - logger.info(f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}") |
118 | | - logger.info(f"backend type: {type(connection.session.backend)}") |
119 | | - |
120 | | - # Close the connection |
121 | | - logger.info("Closing the SEA session...") |
122 | | - connection.close() |
123 | | - logger.info("Successfully closed SEA session") |
124 | | - |
125 | | - except Exception as e: |
126 | | - logger.error(f"Error testing SEA session: {str(e)}") |
127 | | - import traceback |
128 | | - logger.error(traceback.format_exc()) |
| 105 | + logger.error("Please set these variables before running the tests.") |
129 | 106 | sys.exit(1) |
130 | | - |
131 | | - logger.info("SEA session test completed successfully") |
132 | 107 |
|
133 | | -if __name__ == "__main__": |
134 | | - test_sea_session() |
| 108 | + # Run all tests |
| 109 | + results = run_tests() |
| 110 | + |
| 111 | + # Print summary |
| 112 | + print_summary(results) |
| 113 | + |
| 114 | + # Exit with appropriate status code |
| 115 | + all_passed = all(success for _, success in results) |
| 116 | + sys.exit(0 if all_passed else 1) |
0 commit comments