example unit test for missing energy source codes in PUDL allocation

allocation-unit-test.diff

diff --git a/test/unit/analysis/allocate_net_gen_test.py b/test/unit/analysis/allocate_net_gen_test.py
index e3a72a37..d0e50335 100644
--- a/test/unit/analysis/allocate_net_gen_test.py
+++ b/test/unit/analysis/allocate_net_gen_test.py
@@ -1,5 +1,6 @@
 """Unit tests for allocation of net generation."""
 
+from io import StringIO
 from typing import Literal
 
 import numpy as np
@@ -709,29 +710,54 @@ class PudlTablMock:
 
     freq: Literal["AS", "MS"]
 
-    def gens_eia860():
+    def __init__(
+        self,
+        gens_eia860=None,
+        gen_eia923=None,
+        gen_original_eia923=None,
+        generation_fuel_eia923=None,
+        plants_eia860=None,
+        boiler_fuel_eia923=None,
+        boiler_generator_assn_eia860=None,
+        freq="AS",
+    ):
+        self._gens_eia860 = gens_eia860
+        self._gen_eia923 = gen_eia923
+        self._gen_original_eia923 = gen_original_eia923
+        self._generation_fuel_eia923 = generation_fuel_eia923
+        self._plants_eia860 = plants_eia860
+        self._boiler_fuel_eia923 = boiler_fuel_eia923
+        self._boiler_generator_assn_eia860 = boiler_generator_assn_eia860
+
+        self.freq = freq
+
+    def gens_eia860(self):
         """Access to generators_eia860 table."""
-        return GENS_1
+        return self._gens_eia860
+
+    def gen_eia923(self):
+        """Access to generation_eia923 table."""
+        return self._gen_eia923
 
-    def gen_eia923():
+    def gen_original_eia923(self):
         """Access to generation_eia923 table."""
-        return GEN_1
+        return self._gen_original_eia923
 
-    def gf_eia923():
+    def gf_eia923(self):
         """Access to generation_fuel_eia923 table."""
-        return GF_1
+        return self._generation_fuel_eia923
 
-    def plants_eia860():
+    def plants_eia860(self):
         """Access to plants_eia860 table."""
-        return  # PLANTS_1
+        return self._plants_eia860
 
-    def bf_eia923():
+    def bf_eia923(self):
         """Access to boiler_fuel_eia923 table."""
-        return  # BF_1
+        return self._boiler_fuel_eia923
 
-    def bga_eia860():
+    def bga_eia860(self):
         """Access to boiler_generators_assn_eia860 table."""
-        return  # BGA_1
+        return self._boiler_generator_assn_eia860
 
 
 @pytest.mark.xfail(
@@ -855,3 +881,77 @@ def test_allocate_gen_fuel_by_gen_pm_fuel_1():
     ).pipe(apply_pudl_dtypes, group="eia")
     gen_out_1_actual = allocate_net_gen.agg_by_generator(gen_pm_fuel_1_actual)
     pd.testing.assert_frame_equal(gen_out_1_expected, gen_out_1_actual)
+
+
+def test_missing_energy_source():
+
+    gens_eia860 = pd.read_csv(
+        StringIO(
+            """report_date,plant_id_eia,generator_id,prime_mover_code,unit_id_pudl,capacity_mw,fuel_type_count,operational_status,retirement_date,energy_source_code_1,energy_source_code_2,energy_source_code_3,energy_source_code_4,energy_source_code_5,energy_source_code_6,energy_source_code_7,planned_energy_source_code_1,startup_source_code_1,startup_source_code_2,startup_source_code_3,startup_source_code_4
+    2019-01-01,8023,1,ST,1,556.0,1,existing,nan,SUB,BIT,null,null,nan,nan,nan,nan,DFO,nan,nan,nan
+    2019-01-01,8023,2,ST,2,556.0,1,existing,nan,SUB,SUB,BIT,nan,nan,nan,nan,DFO,nan,nan,nan
+    """
+        ),
+    ).pipe(apply_pudl_dtypes, group="eia")
+
+    boiler_fuel_eia923 = pd.read_csv(
+        StringIO(
+            """report_date,plant_id_eia,boiler_id,energy_source_code,prime_mover_code,fuel_consumed_mmbtu
+    2019-01-01,8023,1,DFO,ST,17853.519999999997
+    2019-01-01,8023,1,RC,ST,27681065.276
+    2019-01-01,8023,1,SUB,ST,0.0
+    2019-01-01,8023,2,DFO,ST,17712.999999999996
+    2019-01-01,8023,2,RC,ST,29096935.279
+    2019-01-01,8023,2,SUB,ST,0.0
+    """
+        ),
+    ).pipe(apply_pudl_dtypes, group="eia")
+    # generation_eia923
+    gen_eia923 = pd.read_csv(
+        StringIO(
+            """report_date,plant_id_eia,generator_id,net_generation_mwh
+    2019-01-01,8023,1,2606737.0
+    2019-01-01,8023,2,2759826.0
+    """
+        ),
+    ).pipe(apply_pudl_dtypes, group="eia")
+
+    gen_original_eia923 = gen_eia923
+
+    # boiler_generator_association_eia860
+    boiler_generator_assn_eia860 = pd.read_csv(
+        StringIO(
+            """plant_id_eia,boiler_id,generator_id,report_date
+    8023,1,1,2019-01-01
+    8023,2,2,2019-01-01
+    """
+        ),
+    ).pipe(apply_pudl_dtypes, group="eia")
+
+    generation_fuel_eia923 = pd.read_csv(
+        StringIO(
+            """report_date,plant_id_eia,energy_source_code,prime_mover_code,net_generation_mwh,fuel_consumed_mmbtu,fuel_consumed_for_electricity_mmbtu
+    2019-01-01,8023,DFO,ST,3369.286,35566.0,35566.0
+    2019-01-01,8023,RC,ST,5363193.71,56777578.0,56777578.0
+    2019-01-01,8023,SUB,ST,0.0, 0.0,0.0
+    """
+        ),
+    ).pipe(apply_pudl_dtypes, group="eia")
+
+    mock_pudl_out = PudlTablMock(
+        gens_eia860=gens_eia860,
+        gen_eia923=gen_eia923,
+        gen_original_eia923=gen_original_eia923,
+        generation_fuel_eia923=generation_fuel_eia923,
+        boiler_fuel_eia923=boiler_fuel_eia923,
+        boiler_generator_assn_eia860=boiler_generator_assn_eia860,
+    )
+
+    allocated = allocate_net_gen.allocate_gen_fuel_by_generator_energy_source(
+        mock_pudl_out
+    )
+
+    assert (
+        generation_fuel_eia923.fuel_consumed_mmbtu.sum()
+        == allocated.fuel_consumed_mmbtu.sum()
+    )