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