fix tests

Author: aleksandrachash

.github/workflows/flights_liquibase.yml

@@ -99,6 +99,7 @@ jobs:
           DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }}
           DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN_TST }}
           DATABRICKS_CLUSTER_ID: ${{ secrets.DATABRICKS_CLUSTER_ID }}
+          DBCONNECT_SERVERLESS: "true"
         run: |
           pytest flights/flights-liquibase/tests/ 
       

flights/flights-liquibase/tests/unit_transforms/test_flight_transforms.py

@@ -7,7 +7,13 @@
 def spark_session():
     try:
         from databricks.connect import DatabricksSession
-        return DatabricksSession.builder.profile("adb").getOrCreate()
+        if os.environ.get("DBCONNECT_SERVERLESS", "false").lower() == "true":
+            return DatabricksSession.builder.serverless(True).getOrCreate()
+        else:
+            return DatabricksSession.builder.getOrCreate()   
+    # except (ValueError, RuntimeError):
+    #     from databricks.connect import DatabricksSession
+    #     return DatabricksSession.builder.profile("unit_tests").getOrCreate()    
     except ImportError:
         print("No Databricks Connect, build and return local SparkSession")
         from pyspark.sql import SparkSession