From f99035ebebe781c6f7f5fba4bb634d29847da463 Mon Sep 17 00:00:00 2001
From: Fabian Jakobs <fabian.jakobs@databricks.com>
Date: Wed, 30 Nov 2022 09:29:34 +0100
Subject: [PATCH] Add job to run unit tests

---
 .gitignore                |  2 +-
 .vscode/launch.json       | 16 ++++++++++++++++
 jobs/pytest_databricks.py | 12 ++++++++++++
 tests/spark_test.py       | 18 ++++++++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 .vscode/launch.json
 create mode 100644 jobs/pytest_databricks.py
 create mode 100644 tests/spark_test.py

diff --git a/.gitignore b/.gitignore
index 802ee54..dd48c61 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,8 +1,8 @@
 .databricks/
 .venv/
+.pytest_cache/
 *.pyc
 __pycache__/
-.pytest_cache/
 dist/
 build/
 covid_analysis.egg-info/
\ No newline at end of file
diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 0000000..599bea0
--- /dev/null
+++ b/.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": {}
+    }
+  ]
+}
diff --git a/jobs/pytest_databricks.py b/jobs/pytest_databricks.py
new file mode 100644
index 0000000..4d007e9
--- /dev/null
+++ b/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:])
diff --git a/tests/spark_test.py b/tests/spark_test.py
new file mode 100644
index 0000000..81c1075
--- /dev/null
+++ b/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