Ok, so the best solution as it stands today (for me personally at least) is this:

• Have pyspark ^3.4 installed with the connect extra feature.
• My unit tests then don't have to change at all, as they use the regular spark session created on the fly
• For running the script locally while taking advantage of Databricks, I use the open source Spark Connect and then set the SPARK_REMOTE=sc://${WORKSPACE_INSTANCE_NAME}:443/;token=${PERSONAL_ACCESS_TOKEN};x-databricks-cluster-id=${CLUSTER_ID}