chore: nans to 0s for mapping

agrc · May 30, 2024 · 6ca9a4a · 6ca9a4a
1 parent 3c65f2b
commit 6ca9a4a
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git 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
@@ -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(