diff --git a/src/wmrc/main.py b/src/wmrc/main.py index 7866027..c5476ef 100644 --- a/src/wmrc/main.py +++ b/src/wmrc/main.py @@ -368,6 +368,7 @@ def _county_summaries(records: helpers.SalesForceRecords) -> pd.DataFrame: index={name: name.replace("__c", "").replace("_", " ") for name in county_df.index}, inplace=True ) county_df["data_year"] = county_df["data_year"].apply(helpers.convert_to_int) + county_df.fillna(0, inplace=True) return county_df diff --git a/tests/test_wmrc.py b/tests/test_wmrc.py index 917dd0a..dc62bbe 100644 --- a/tests/test_wmrc.py +++ b/tests/test_wmrc.py @@ -1,3 +1,4 @@ +import numpy as np import pandas as pd from wmrc import main @@ -102,6 +103,41 @@ def test_county_summaries_happy_path(self, mocker): pd.testing.assert_frame_equal(result_df, test_df) + def test_county_summaries_replace_nan_with_0(self, mocker): + records_mock = mocker.Mock() + summaries_df = pd.DataFrame( + { + "recycled": [1, 2, 3, 4], + "landfilled": [5, 6, np.nan, 8], + "total": [9, 10, np.nan, 12], + }, + index=pd.MultiIndex.from_tuples( + [ + ("2022", "Box Elder County"), + ("2022", "Out of State"), + ("2023", "Box Elder County"), + ("2023", "Out of State"), + ], + names=["year", "county"], + ), + ) + records_mock.df.groupby.return_value.apply.return_value = summaries_df + + result_df = main.Skid._county_summaries(records_mock) + + test_df = pd.DataFrame( + { + "data_year": [2022, 2022, 2023, 2023], + "recycled": [1, 2, 3, 4], + "landfilled": [5, 6, 0.0, 8], + "total": [9, 10, 0.0, 12], + }, + index=["Box Elder County", "Out of State", "Box Elder County", "Out of State"], + ) + test_df.index.name = "name" + + pd.testing.assert_frame_equal(result_df, test_df) + def test_contamination_rates_by_tonnage_happy_path(self, mocker): records = mocker.Mock() records.df = pd.DataFrame(