Merge pull request #11 from fjakobs/fixes

Add support for running unit tests from VSCode

Author: fjakobs

.gitignore

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

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+  "python.testing.pytestArgs": ["."],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true
+}

jobs/covid_trends_job.py

@@ -2,7 +2,6 @@
 Python Spark job that imports the latest COVID-19 hospitalization data
 '''
 import sys
-import urllib.request
 import pandas as pd
 from pyspark.sql import SparkSession
 
@@ -13,27 +12,26 @@
 # check if job is running in production mode
 is_prod = len(sys.argv) >= 2 and sys.argv[1] == "--prod"
 
-urllib.request.urlretrieve("https://raw.githubusercontent.com/owid/covid-19-data/master/public/data/hospitalizations/covid-hospitalizations.csv", "/tmp/covid-hospitalizations.csv")
- 
 # read from /tmp, subset for USA, pivot and fill missing values
-df = pd.read_csv("/tmp/covid-hospitalizations.csv")
-df = filter_country(df, country='DZA')
-df = pivot_and_clean(df, fillna=0)  
+df = pd.read_csv(
+    "https://raw.githubusercontent.com/owid/covid-19-data/master/public/data/hospitalizations/covid-hospitalizations.csv")
+
+df = filter_country(df)
+df = pivot_and_clean(df, fillna=0)
 df = clean_spark_cols(df)
 df = index_to_col(df, colname='date')
 
 # Convert from Pandas to a pyspark sql DataFrame.
 df = spark.createDataFrame(df)
 
-print("Covid data successfully imported.")
-
 # only write table in production mode
 if is_prod:
     # Write to Delta Lake
     df.write.mode('overwrite').saveAsTable('covid_stats')
+    print("Covid data successfully imported.")
 
-    # display sample data
+# display sample data
+if is_prod:
     spark.sql('select * from covid_stats').show(10)
-
-
-
+else:
+    df.show(10)

tests/__init__.py

@@ -0,0 +1,2 @@
+import sys
+sys.path.append('.')

tests/transforms_test.py

@@ -1,11 +1,7 @@
 # Test each of the transform functions.
 import pytest
-from textwrap import fill
-import os
 import pandas as pd
-import numpy as np
 from covid_analysis.transforms import *
-from pyspark.sql import SparkSession
 
 
 @pytest.fixture
@@ -30,7 +26,7 @@ def colnames_df() -> pd.DataFrame:
         ],
     )
     return df
-  
+
 
 # Make sure the filter works as expected.
 def test_filter(raw_input_df):