from databricks.connect import DatabricksSession
from pyspark.sql import SparkSession
from {{.project_name}} import main

# Create a new Databricks Connect session. If this fails,
# check that you have configured Databricks Connect correctly.
# See https://docs.databricks.com/dev-tools/databricks-connect.html.
{{/*
  The below works around a problematic error message from Databricks Connect.
  The standard SparkSession is supported in all configurations (workspace, IDE,
  all runtime versions, CLI). But on the CLI it currently gives a confusing
  error message if SPARK_REMOTE is not set. We can't directly use
  DatabricksSession.builder in main.py, so we're re-assigning it here so
  everything works out of the box, even for CLI users who don't set SPARK_REMOTE.
*/}}
SparkSession.builder = DatabricksSession.builder
SparkSession.builder.getOrCreate()

def test_main():
    taxis = main.get_taxis()
    assert taxis.count() > 5
