Remove sample.after_lp

tests/components/test_objective.py

@@ -19,19 +19,12 @@ from miplearn.solvers.pyomo.gurobi import GurobiPyomoSolver
 @pytest.fixture
 def sample() -> Sample:
     sample = Sample(
-        after_load=Features(
-            instance=InstanceFeatures(),
-        ),
-        after_lp=Features(
-            lp_solve=LPSolveStats(),
-        ),
         data={
             "mip_lower_bound": 1.0,
             "mip_upper_bound": 2.0,
+            "lp_instance_features": [1.0, 2.0, 3.0],
         },
     )
-    sample.after_load.instance.to_list = Mock(return_value=[1.0, 2.0])  # type: ignore
-    sample.after_lp.lp_solve.to_list = Mock(return_value=[3.0])  # type: ignore
     return sample
 
 

tests/components/test_primal.py

@@ -33,12 +33,23 @@ def sample() -> Sample:
                 categories=["default", None, "default", "default"],
             ),
         ),
-        after_lp=Features(
-            variables=VariableFeatures(),
-        ),
         data={
             "var_names": ["x[0]", "x[1]", "x[2]", "x[3]"],
+            "var_categories": ["default", None, "default", "default"],
             "mip_var_values": [0.0, 1.0, 1.0, 0.0],
+            "instance_features_user": [5.0],
+            "var_features": [
+                [0.0, 0.0],
+                None,
+                [1.0, 0.0],
+                [1.0, 1.0],
+            ],
+            "lp_var_features": [
+                [0.0, 0.0, 2.0, 2.0],
+                None,
+                [1.0, 0.0, 3.0, 2.0],
+                [1.0, 1.0, 3.0, 3.0],
+            ],
         },
     )
     sample.after_load.instance.to_list = Mock(return_value=[5.0])  # type: ignore
@@ -50,14 +61,6 @@ def sample() -> Sample:
             [1.0, 1.0],
         ][i]
     )
-    sample.after_lp.variables.to_list = Mock(  # type:ignore
-        side_effect=lambda i: [
-            [2.0, 2.0],
-            None,
-            [3.0, 2.0],
-            [3.0, 3.0],
-        ][i]
-    )
     return sample
 
 

tests/components/test_static_lazy.py

@@ -44,24 +44,28 @@ def sample() -> Sample:
                 lazy=[True, True, True, True, False],
             ),
         ),
-        after_lp=Features(
-            instance=InstanceFeatures(),
-            constraints=ConstraintFeatures(names=["c1", "c2", "c3", "c4", "c5"]),
-        ),
         data={
+            "constr_categories": [
+                "type-a",
+                "type-a",
+                "type-a",
+                "type-b",
+                "type-b",
+            ],
+            "constr_lazy": [True, True, True, True, False],
+            "constr_names": ["c1", "c2", "c3", "c4", "c5"],
+            "instance_features_user": [5.0],
             "lazy_enforced": {"c1", "c2", "c4"},
+            "lp_constr_features": [
+                [1.0, 1.0],
+                [1.0, 2.0],
+                [1.0, 3.0],
+                [1.0, 4.0, 0.0],
+                None,
+            ],
+            "static_lazy_count": 4,
         },
     )
-    sample.after_lp.instance.to_list = Mock(return_value=[5.0])  # type: ignore
-    sample.after_lp.constraints.to_list = Mock(  # type: ignore
-        side_effect=lambda idx: {
-            0: [1.0, 1.0],
-            1: [1.0, 2.0],
-            2: [1.0, 3.0],
-            3: [1.0, 4.0, 0.0],
-            4: None,
-        }[idx]
-    )
     return sample