1
+ import shutil
2
+ from datetime import datetime
1
3
from os import listdir
2
4
from os .path import join
3
5
from itertools import product
4
6
5
7
import pandas as pd
6
8
9
+ from conftest import TEST_DIR
7
10
8
11
class TestRun :
9
- def test_output_files_exist (self , run_as_module ):
10
- csv_files = listdir ("receiving" )
12
+ def test_output_files_exist (self ):
13
+ output_files = listdir ("receiving" )
14
+ smoothed_files = sorted (list (set ([file for file in output_files if "smoothed" in file ])))
15
+ raw_files = sorted (list (set ([file for file in output_files if "raw" in file ])))
16
+ csv_files = {"raw" : raw_files , "smoothed" : smoothed_files }
11
17
12
- dates = [
18
+ expected_smoothed_dates = [
13
19
"20200801" ,
14
20
"20200802" ,
15
21
"20200803" ,
@@ -22,22 +28,41 @@ def test_output_files_exist(self, run_as_module):
22
28
"20200810" ,
23
29
"20200811"
24
30
]
31
+ expected_raw_dates = [
32
+ '20200726' ,
33
+ '20200727' ,
34
+ '20200728' ,
35
+ '20200729' ,
36
+ '20200730' ,
37
+ '20200731' ,
38
+ ] + expected_smoothed_dates
39
+
40
+ dates = {
41
+ "raw" : expected_raw_dates ,
42
+ "smoothed" : expected_smoothed_dates ,
43
+ }
44
+
25
45
geos = ["county" , "state" , "hhs" , "nation" ]
26
46
metrics = ["s01" , "s02" , "s03" ,
27
- #"s04",
47
+ # "s04",
28
48
"s04" , "s05" ,
29
- #"s07",
49
+ # "s07",
30
50
"s06" ,
31
- #"s09", "s10",
51
+ # "s09", "s10",
32
52
"scontrol" ]
33
53
smoother = ["raw" , "smoothed" ]
34
54
35
- expected_files = []
36
- for date , geo , metric , smoother in product (dates , geos , metrics , smoother ):
37
- nf = "_" .join ([date , geo , metric , smoother , "research" ]) + ".csv"
38
- expected_files .append (nf )
55
+ for smther in smoother :
56
+ expected_files = []
57
+
58
+ for date , geo , metric in product (dates [smther ], geos , metrics ):
59
+ nf = "_" .join ([date , geo , metric , smther , "search" ]) + ".csv"
60
+ expected_files .append (nf )
61
+
62
+ csv_dates = list (set ([datetime .strptime (f .split ('_' )[0 ], "%Y%m%d" ) for f in csv_files [smther ] if smther in f ]))
63
+ assert set (csv_files [smther ]).issuperset (set (expected_files ))
39
64
40
- set ( csv_files ) == set ( expected_files )
65
+ shutil . rmtree ( f" { TEST_DIR } /receiving/" )
41
66
42
67
def test_output_file_format (self , run_as_module ):
43
68
df = pd .read_csv (
0 commit comments