From 0b961e34ebf55041509124a1c1fc402b43ab60dc Mon Sep 17 00:00:00 2001
From: John Vivian <jtvivian@gmail.com>
Date: Fri, 20 Oct 2023 20:40:56 -0700
Subject: [PATCH] Update tests

---
 coverage.xml             |  2 +-
 tests/test_dfm.py        | 11 ++++++-----
 tests/test_processing.py |  8 ++++----
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/coverage.xml b/coverage.xml
index 2bc0a41..eeb38d0 100644
--- a/coverage.xml
+++ b/coverage.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" ?>
-<coverage version="7.3.1" timestamp="1697579914664" lines-valid="115" lines-covered="115" line-rate="1" branches-valid="42" branches-covered="42" branch-rate="1" complexity="0">
+<coverage version="7.3.1" timestamp="1697606544233" lines-valid="115" lines-covered="115" line-rate="1" branches-valid="42" branches-covered="42" branch-rate="1" complexity="0">
 	<!-- Generated by coverage.py: https://coverage.readthedocs.io/en/7.3.1 -->
 	<!-- Based on https://raw.githubusercontent.com/cobertura/web/master/htdocs/xml/coverage-04.dtd -->
 	<sources>
diff --git a/tests/test_dfm.py b/tests/test_dfm.py
index 608b437..4003077 100644
--- a/tests/test_dfm.py
+++ b/tests/test_dfm.py
@@ -1,14 +1,15 @@
 import os
+import shutil
 from pathlib import Path
 
 from covid19_drdfm.dfm import run_model
 from covid19_drdfm.processing import get_df
 
+
 # TODO: output should go in a directory instead of dumping shit everywhere
 def test_run_model():
     df = get_df()
-    run_model(df, "NY", Path("./"))
-    assert Path("./model.csv").exists()
-    assert Path("./results.csv").exists()
-    os.remove("./model.csv")
-    os.remove("./results.csv")
+    run_model(df, "NY", Path("./testdir"))
+    assert Path("./testdir/model.csv").exists()
+    assert Path("./testdir/results.csv").exists()
+    shutil.rmtree("./testdir")
diff --git a/tests/test_processing.py b/tests/test_processing.py
index 8309344..771076b 100644
--- a/tests/test_processing.py
+++ b/tests/test_processing.py
@@ -1,16 +1,16 @@
-import pandas as pd
-import pytest
 from functools import reduce
 
+import pandas as pd
+import pytest
 
 from covid19_drdfm.processing import (
+    DATA_DIR,
+    ROOT_DIR,
     add_datetime,
     adjust_inflation,
     adjust_pandemic_response,
     get_df,
     get_govt_fund_dist,
-    DATA_DIR,
-    ROOT_DIR,
 )