fix disease tests

ihmeuw · Oct 24, 2023 · e753aed · e753aed
1 parent 4e4f62d
commit e753aed
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/tests/disease/test_disease.py b/tests/disease/test_disease.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from vivarium import InteractiveContext
+from vivarium import InteractiveContext, Component
 from vivarium.framework.state_machine import Transition
 from vivarium.framework.utilities import from_yearly
 from vivarium.testing_utilities import TestPopulation, build_table, metadata
@@ -358,10 +358,7 @@ def test_risk_deletion(base_config, base_plugins, disease):
 
     model = DiseaseModel(disease, initial_state=healthy, states=[healthy, sick])
 
-    class PafModifier:
-        @property
-        def name(self):
-            return "paf_modifier"
+    class PafModifier(Component):
 
         def setup(self, builder):
             builder.value.register_value_modifier(
@@ -370,7 +367,7 @@ def setup(self, builder):
                     build_table(paf, year_start, year_end),
                     key_columns=("sex",),
                     parameter_columns=["age", "year"],
-                    value_columns=None,
+                    value_columns=(),
                 ),
             )