Update precondition tests

Author: cthoyt

src/y0/controls.py

@@ -12,8 +12,6 @@
 
 
 def _control_precondition(graph: NxMixedGraph, query: Probability, variable: Variable):
-    if missing := query.get_variables().difference(graph.nodes()):
-        raise ValueError(f"Query variables missing: {missing}")
     if variable not in graph.nodes():
         raise ValueError(f"Test variable missing: {variable}")
     # TODO does this need to be extended to check that the
@@ -29,7 +27,6 @@ def is_good_control(graph: NxMixedGraph, query: Probability, variable: Variable)
     :return: If the variable is a good control
     """
     _control_precondition(graph, query, variable)
-
     raise NotImplementedError
 
 
@@ -46,5 +43,4 @@ def is_bad_control(graph: NxMixedGraph, query: Probability, variable: Variable)
     :return: If the variable is a bad control
     """
     _control_precondition(graph, query, variable)
-
     raise NotImplementedError

tests/test_controls.py

@@ -5,7 +5,7 @@
 import unittest
 
 from y0.controls import is_bad_control, is_good_control
-from y0.dsl import U1, U2, M, P, U, W, X, Y, Z
+from y0.dsl import U1, U2, A, M, P, U, W, X, Y, Z
 from y0.graph import NxMixedGraph
 
 model_1 = NxMixedGraph.from_edges(directed=[(Z, X), (Z, Y), (X, Y)])
@@ -71,6 +71,13 @@
 class TestControls(unittest.TestCase):
     """Test case for good, bad, and neutral controls."""
 
+    def test_preconditions(self):
+        """Test the preconditions are checked properly for good controls."""
+        for func in is_good_control, is_bad_control:
+            with self.subTest(name=func.__name__):
+                with self.assertRaises(ValueError):
+                    func(model_1, P(Y @ X), A)
+
     def test_good_controls(self):
         """Test good controls."""
         for model in good_test_models: