def test_line_plot(self, jhu_data, population_data, country):
warnings.simplefilter("ignore", category=UserWarning)
# Setting
dhl = DataHandler(jhu_data, population_data, country)
dhl.init_records()
# Interactive / script mode
assert not dhl.interactive
with pytest.raises(NotInteractiveError):
dhl.interactive = True
dhl.interactive = False
# Change colors in plotting
dhl.records(
variables=["Confirmed", "Infected", "Fatal", "Recovered"],
color_dict={"Confirmed": "blue", "Infected": "orange", "Fatal": "red", "Recovered": "green"})
def test_records(self, jhu_data, population_data, country):
warnings.simplefilter("ignore", category=UserWarning)
# Setting
dhl = DataHandler(jhu_data, population_data, country)
dhl.init_records()
dhl.first_date = "01Apr2020"
dhl.last_date = "01Aug2020"
# Test
df = dhl.records(variables=None, show_figure=False)
assert isinstance(df, pd.DataFrame)
assert set(df.columns) == set([Term.DATE, Term.CI, Term.F, Term.R])
dates = df[Term.DATE]
assert dates.min() == Term.date_obj(dhl.first_date)
assert dates.max() == Term.date_obj(dhl.last_date)
df2 = dhl.records(variables=["Susceptible"], show_figure=True)
assert set(df2.columns) == set([Term.DATE, Term.S])
def test_records(self, jhu_data, population_data, country, japan_data,
oxcgrt_data, pcr_data, vaccine_data, main, extras, past,
future):
dhl = DataHandler(country=country, province=None)
# Combination of arguments
if (not main and not extras) or (not past and not future):
with pytest.raises(ValueError):
dhl.records(main=main, extras=extras, past=past, future=future)
return
# Register main datasets
if main:
with pytest.raises(NotRegisteredMainError):
dhl.records(main=main, extras=extras, past=past, future=future)
dhl.register(jhu_data=jhu_data, population_data=population_data)
dhl.timepoints(first_date="01Apr2020",
last_date="01Sep2020",
today="01Jun2020")
# Register extra datasets
if extras:
with pytest.raises(NotRegisteredExtraError):
dhl.records(main=main, extras=extras, past=past, future=future)
dhl.register(extras=[japan_data, oxcgrt_data, pcr_data, vaccine_data])
# Get records
df = dhl.records(main=main, extras=extras, past=past, future=future)
# Check the start/end date of the records
if past and future:
sta, end = "01Apr2020", "01Sep2020"
elif past:
sta, end = "01Apr2020", "01Jun2020"
else:
sta, end = "02Jun2020", "01Sep2020"
assert df[Term.DATE].min().strftime(Term.DATE_FORMAT) == sta
assert df[Term.DATE].max().strftime(Term.DATE_FORMAT) == end
