test: fix unittests

Author: TomeHirata

dte_adj/__init__.py

@@ -355,12 +355,12 @@ def fit(
             SimpleDistributionEstimator: The fitted estimator.
         """
         if confoundings.shape[0] != treatment_arms.shape[0]:
-            raise RuntimeError(
+            raise ValueError(
                 "The shape of confounding and treatment_arm should be same"
             )
 
         if confoundings.shape[0] != outcomes.shape[0]:
-            raise RuntimeError("The shape of confounding and outcome should be same")
+            raise ValueError("The shape of confounding and outcome should be same")
 
         self.confoundings = confoundings
         self.treatment_arms = treatment_arms
@@ -379,7 +379,7 @@ def predict(self, treatment_arms: np.ndarray, locations: np.ndarray) -> np.ndarr
             np.ndarray: Estimated cumulative distribution values for the input.
         """
         if self.outcomes is None:
-            raise RuntimeError(
+            raise ValueError(
                 "This estimator has not been trained yet. Please call fit first"
             )
 
@@ -461,12 +461,12 @@ def fit(
             AdjustedDistributionEstimator: The fitted estimator.
         """
         if confoundings.shape[0] != treatment_arms.shape[0]:
-            raise RuntimeError(
+            raise ValueError(
                 "The shape of confounding and treatment_arm should be same"
             )
 
         if confoundings.shape[0] != outcomes.shape[0]:
-            raise RuntimeError("The shape of confounding and outcome should be same")
+            raise ValueError("The shape of confounding and outcome should be same")
 
         self.confoundings = confoundings
         self.treatment_arms = treatment_arms
@@ -485,7 +485,7 @@ def predict(self, treatment_arms: np.ndarray, locations: np.ndarray) -> np.ndarr
             np.ndarray: Estimated cumulative distribution values for the input.
         """
         if self.outcomes is None:
-            raise RuntimeError(
+            raise ValueError(
                 "This estimator has not been trained yet. Please call fit first"
             )
 

dte_adj/util.py

@@ -105,7 +105,7 @@ def compute_confidence_intervals(
 
         return vec_dte_lower_simple, vec_dte_upper_simple
     else:
-        raise RuntimeError(f"Invalid variance type was speficied: {variance_type}")
+        raise ValueError(f"Invalid variance type was speficied: {variance_type}")
 
 
 def find_le(array: np.ndarray, threshold):

tests/test_adjusted_estimator.py

@@ -18,7 +18,7 @@ def test_prediction_fail_before_fit(self):
         subject = AdjustedDistributionEstimator(base_model)
 
         # Act, Assert
-        with self.assertRaises(RuntimeError) as cm:
+        with self.assertRaises(ValueError) as cm:
             subject.predict(D, Y)
         self.assertEqual(
             str(cm.exception),
@@ -35,7 +35,7 @@ def test_fit_fail_invalid_input(self):
         subject = AdjustedDistributionEstimator(base_model)
 
         # Act, Assert
-        with self.assertRaises(RuntimeError) as cm:
+        with self.assertRaises(ValueError) as cm:
             subject.fit(X, D, Y)
         self.assertEqual(
             str(cm.exception),

tests/test_plot.py

@@ -9,7 +9,7 @@ class TestPlot(unittest.TestCase):
     def test_plot(self, mock_plt):
         # Arrange
         x_values = np.array([1, 2, 3, 4, 5])
-        y_values = np.array([1, 2, 3, 4, 5])
+        means = np.array([1, 2, 3, 4, 5])
         upper_bands = np.array([2, 3, 4, 5, 6])
         lower_bands = np.array([0, 1, 2, 3, 4])
         mock_ax = MagicMock()
@@ -18,12 +18,13 @@ def test_plot(self, mock_plt):
         # Act
         result_ax = plot(
             x_values,
-            y_values,
-            upper_bands,
+            means,
             lower_bands,
+            upper_bands,
             title="Test Title",
             xlabel="X Axis",
             ylabel="Y Axis",
+            chart_type="line",
         )
 
         # Assert
@@ -34,7 +35,7 @@ def test_plot(self, mock_plt):
         plot_args, plot_kwargs = plot_call
         x_values_arg, y_values_arg = plot_args
         self.assertTrue(np.array_equal(x_values_arg, x_values))
-        self.assertTrue(np.array_equal(y_values_arg, y_values))
+        self.assertTrue(np.array_equal(y_values_arg, means))
         fill_between_args, fill_between_kwargs = fill_between_call
         x_fill, lower_fill, upper_fill = fill_between_args
         self.assertTrue(np.array_equal(x_fill, x_values_arg))
@@ -46,8 +47,30 @@ def test_plot(self, mock_plt):
         mock_ax.set_title.assert_called_once_with("Test Title")
         mock_ax.set_xlabel.assert_called_once_with("X Axis")
         mock_ax.set_ylabel.assert_called_once_with("Y Axis")
-        mock_ax.legend.assert_called_once()
 
+    def test_plot_fail_unknown_chart_type(self):
+        # Arrange
+        x_values = np.array([1, 2, 3, 4, 5])
+        means = np.array([1, 2, 3, 4, 5])
+        upper_bands = np.array([2, 3, 4, 5, 6])
+        lower_bands = np.array([0, 1, 2, 3, 4])
+
+        # Act, Assert
+        with self.assertRaises(ValueError) as cm:
+          plot(
+              x_values,
+              means,
+              lower_bands,
+              upper_bands,
+              title="Test Title",
+              xlabel="X Axis",
+              ylabel="Y Axis",
+              chart_type="other",
+          )
+        self.assertEqual(
+            str(cm.exception),
+            "Chart type other is not supported",
+        )
 
 if __name__ == "__main__":
     unittest.main()

tests/test_simple_estimator.py

@@ -31,7 +31,7 @@ def test_prediction_fail_before_fit(self):
         subject = SimpleDistributionEstimator()
 
         # Act, Assert
-        with self.assertRaises(RuntimeError) as cm:
+        with self.assertRaises(ValueError) as cm:
             subject.predict(D, Y)
         self.assertEqual(
             str(cm.exception),
@@ -47,7 +47,7 @@ def test_fit_fail_invalid_input(self):
         subject = SimpleDistributionEstimator()
 
         # Act, Assert
-        with self.assertRaises(RuntimeError) as cm:
+        with self.assertRaises(ValueError) as cm:
             subject.fit(X, D, Y)
         self.assertEqual(
             str(cm.exception),