6
6
import pandas as pd
7
7
from delphi_utils import create_export_csv
8
8
9
+ def _clean_directory (directory ):
10
+ """Clean files out of a directory."""
11
+ for fname in listdir (directory ):
12
+ if fname .startswith ("." ):
13
+ continue
14
+ remove (join (directory , fname ))
15
+
16
+
17
+ def _non_ignored_files_set (directory ):
18
+ """List all files in a directory not preceded by a '.' and store them in a set."""
19
+ out = set ()
20
+ for fname in listdir (directory ):
21
+ if fname .startswith ("." ):
22
+ continue
23
+ out .add (fname )
24
+ return out
25
+
9
26
10
27
class TestExport :
28
+ """Tests for exporting CSVs."""
11
29
# List of times for data points.
12
30
TIMES = [
13
31
datetime .strptime (x , "%Y-%m-%d" )
14
32
for x in ["2020-02-15" , "2020-02-15" , "2020-03-01" , "2020-03-15" ]
15
33
]
16
34
17
- # A sample data frame
35
+ # A sample data frame.
18
36
DF = pd .DataFrame (
19
37
{
20
38
"geo_id" : ["51093" , "51175" , "51175" , "51620" ],
@@ -25,23 +43,25 @@ class TestExport:
25
43
}
26
44
)
27
45
46
+ # Directory in which to store tests.
47
+ TEST_DIR = "test_dir"
48
+
28
49
def test_export_with_metric (self ):
29
50
"""Test that exporting CSVs with the `metrics` argument yields the correct files."""
30
51
31
52
# Clean receiving directory
32
- for fname in listdir ("test_dir" ):
33
- remove (join ("test_dir" , fname ))
53
+ _clean_directory (self .TEST_DIR )
34
54
35
55
create_export_csv (
36
56
df = self .DF ,
37
57
start_date = datetime .strptime ("2020-02-15" , "%Y-%m-%d" ),
38
- export_dir = "test_dir" ,
58
+ export_dir = self . TEST_DIR ,
39
59
metric = "deaths" ,
40
60
geo_res = "county" ,
41
61
sensor = "test" ,
42
62
)
43
63
44
- assert set ( listdir ( "test_dir" ) ) == set (
64
+ assert _non_ignored_files_set ( self . TEST_DIR ) == set (
45
65
[
46
66
"20200215_county_deaths_test.csv" ,
47
67
"20200301_county_deaths_test.csv" ,
@@ -53,18 +73,17 @@ def test_export_without_metric(self):
53
73
"""Test that exporting CSVs without the `metrics` argument yields the correct files."""
54
74
55
75
# Clean receiving directory
56
- for fname in listdir ("test_dir" ):
57
- remove (join ("test_dir" , fname ))
76
+ _clean_directory (self .TEST_DIR )
58
77
59
78
create_export_csv (
60
79
df = self .DF ,
61
80
start_date = datetime .strptime ("2020-02-15" , "%Y-%m-%d" ),
62
- export_dir = "test_dir" ,
81
+ export_dir = self . TEST_DIR ,
63
82
geo_res = "county" ,
64
83
sensor = "test" ,
65
84
)
66
85
67
- assert set ( listdir ( "test_dir" ) ) == set (
86
+ assert _non_ignored_files_set ( self . TEST_DIR ) == set (
68
87
[
69
88
"20200215_county_test.csv" ,
70
89
"20200301_county_test.csv" ,
@@ -76,18 +95,17 @@ def test_export_with_limiting_start_date(self):
76
95
"""Test that the `start_date` prevents earlier dates from being exported."""
77
96
78
97
# Clean receiving directory
79
- for fname in listdir ("test_dir" ):
80
- remove (join ("test_dir" , fname ))
98
+ _clean_directory (self .TEST_DIR )
81
99
82
100
create_export_csv (
83
101
df = self .DF ,
84
102
start_date = datetime .strptime ("2020-02-20" , "%Y-%m-%d" ),
85
- export_dir = "test_dir" ,
103
+ export_dir = self . TEST_DIR ,
86
104
geo_res = "county" ,
87
105
sensor = "test" ,
88
106
)
89
107
90
- assert set ( listdir ( "test_dir" ) ) == set (
108
+ assert _non_ignored_files_set ( self . TEST_DIR ) == set (
91
109
[
92
110
"20200301_county_test.csv" ,
93
111
"20200315_county_test.csv" ,
@@ -98,17 +116,16 @@ def test_export_with_no_start_date(self):
98
116
"""Test that omitting the `start_date` exports all dates."""
99
117
100
118
# Clean receiving directory
101
- for fname in listdir ("test_dir" ):
102
- remove (join ("test_dir" , fname ))
119
+ _clean_directory (self .TEST_DIR )
103
120
104
121
create_export_csv (
105
122
df = self .DF ,
106
- export_dir = "test_dir" ,
123
+ export_dir = self . TEST_DIR ,
107
124
geo_res = "state" ,
108
125
sensor = "test" ,
109
126
)
110
127
111
- assert set ( listdir ( "test_dir" ) ) == set (
128
+ assert _non_ignored_files_set ( self . TEST_DIR ) == set (
112
129
[
113
130
"20200215_state_test.csv" ,
114
131
"20200301_state_test.csv" ,
0 commit comments