Add job to run unit tests

.gitignore

@@ -1,8 +1,8 @@
 .databricks/
 .venv/
+.pytest_cache/
 *.pyc
 __pycache__/
-.pytest_cache/
 dist/
 build/
 covid_analysis.egg-info/

.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "type": "databricks",
+      "request": "launch",
+      "name": "Unit Tests (on Databricks)",
+      "program": "${workspaceFolder}/jobs/pytest_databricks.py",
+      "args": ["./tests"],
+      "env": {}
+    }
+  ]
+}

jobs/pytest_databricks.py

@@ -0,0 +1,12 @@
+import pytest
+import os
+import sys
+
+# Run all tests in the repository root.
+repo_root = os.path.dirname(os.getcwd())
+os.chdir(repo_root)
+
+# Skip writing pyc files on a readonly filesystem.
+sys.dont_write_bytecode = True
+
+_ = pytest.main(sys.argv[1:])

tests/spark_test.py

@@ -0,0 +1,18 @@
+from pyspark.sql import SparkSession
+import pytest
+
+
+@pytest.fixture
+def spark() -> SparkSession:
+    """
+    Create a spark session. Unit tests don't have access to the spark global
+    """
+    return SparkSession.builder.getOrCreate()
+
+
+def test_spark(spark):
+    """
+    Example test that needs to run on the cluster to work
+    """
+    data = spark.sql("select 1").collect()
+    assert data[0][0] == 1