Revert test which was broken in last commit

README.md

@@ -34,8 +34,8 @@ Now download a3fe, install, and test:
 - In the calculation base directory, run the following python code, either through ipython or as a python script (you will likely want to run the script with `nohup`or use ipython through tmux to ensure that the calculation is not killed when you lose connection)
 
 ```python
-import a3fe as ee 
-calc = ee.Calculation()
+import a3fe as a3 
+calc = a3.Calculation()
 calc.setup()
 calc.get_optimal_lam_vals()
 calc.run()

a3fe/tests/fixtures.py

@@ -8,7 +8,7 @@
 import BioSimSpace.Sandpit.Exscientia as BSS
 import pytest
 
-import a3fe as ee
+import a3fe as a3
 
 
 @pytest.fixture(scope="session")
@@ -17,9 +17,9 @@ def restrain_stage():
     with TemporaryDirectory() as dirname:
         # Copy the input files to the temporary directory
         subprocess.run(["cp", "-r", "a3fe/data/example_restraint_stage/", dirname])
-        stage = ee.Stage(
+        stage = a3.Stage(
             base_dir=os.path.join(dirname, "example_restraint_stage"),
-            stage_type=ee.enums.StageType.RESTRAIN,
+            stage_type=a3.enums.StageType.RESTRAIN,
         )
         # Set the relative simuation cost to 1
         stage.recursively_set_attr("relative_simulation_cost", 1, force=True)
@@ -31,25 +31,25 @@ def restrain_stage():
 @pytest.fixture(scope="session")
 def restrain_stage_iterator(restrain_stage):
     """Create a simulation runner iterator object with analysis data to use in tests"""
-    restrain_stage_iterator = ee.run._utils.SimulationRunnerIterator(
+    restrain_stage_iterator = a3.run._utils.SimulationRunnerIterator(
         base_dirs=[restrain_stage.base_dir, restrain_stage.base_dir],
-        subclass=ee.Stage,
-        stage_type=ee.StageType.RESTRAIN,
+        subclass=a3.Stage,
+        stage_type=a3.StageType.RESTRAIN,
     )
     yield restrain_stage_iterator
 
 
 @pytest.fixture(scope="session")
 def restrain_stage_grad_data(restrain_stage):
     """Create a gradient data object with analysis data"""
-    yield ee.analyse.GradientData(restrain_stage.lam_windows, equilibrated=True)
+    yield a3.analyse.GradientData(restrain_stage.lam_windows, equilibrated=True)
 
 
 @pytest.fixture(scope="session")
 def calc():
     """Create a calculation object to use in tests"""
     with TemporaryDirectory() as dirname:
-        calc = ee.Calculation(
+        calc = a3.Calculation(
             base_dir=dirname,
             input_dir="a3fe/data/example_run_dir/input",
             ensemble_size=6,

a3fe/tests/test_run.py

@@ -31,7 +31,7 @@ def test_calculation_loading(calc):
     )
     assert calc.output_dir == os.path.join(calc.base_dir, "output")
     assert calc.setup_complete == False
-    assert calc.prep_stage == a3.run.enums.PreparationStage.PARAMETERISED
+    assert calc.prep_stage.name == a3.run.enums.PreparationStage.PARAMETERISED.name
     assert calc.stream_log_level == logging.INFO
     # Check that pickle file exists
     assert os.path.exists(os.path.join(calc.base_dir, "Calculation.pkl"))
@@ -63,7 +63,7 @@ def test_calculation_reloading(calc):
     )
     assert calc2.output_dir == os.path.join(calc.base_dir, "output")
     assert calc2.setup_complete == False
-    assert calc2.prep_stage == a3.run.enums.PreparationStage.PARAMETERISED
+    assert calc2.prep_stage.name == a3.run.enums.PreparationStage.PARAMETERISED.name
     assert calc2.stream_log_level == logging.INFO
 
 
@@ -93,17 +93,17 @@ def test_set_and_get_attributes(restrain_stage):
         == 5
     )
     # Check it fails if the attribute doesn't exist
-    restrain_stage.recursively_set_attr("ensemble_size", 7)
-    attr_dict = restrain_stage.recursively_get_attr("ensemble_size")
-    assert attr_dict["ensemble_size"] == None
+    restrain_stage.recursively_set_attr("ensemble_sizee", 7)
+    attr_dict = restrain_stage.recursively_get_attr("ensemble_sizee")
+    assert attr_dict["ensemble_sizee"] == None
     # Check that we can force it to set the attribute
-    restrain_stage.recursively_set_attr("ensemble_size", 7, force=True)
-    attr_dict = restrain_stage.recursively_get_attr("ensemble_size")
-    assert attr_dict["ensemble_size"] == 7
+    restrain_stage.recursively_set_attr("ensemble_sizee", 7, force=True)
+    attr_dict = restrain_stage.recursively_get_attr("ensemble_sizee")
+    assert attr_dict["ensemble_sizee"] == 7
     # Change the ensemble size attribute
-    restrain_stage.recursively_set_attr("ensemble_size", 7)
-    attr_dict = restrain_stage.recursively_get_attr("ensemble_size")
-    assert attr_dict["ensemble_size"] == 7
+    restrain_stage.recursively_set_attr("ensemble_sizee", 7)
+    attr_dict = restrain_stage.recursively_get_attr("ensemble_sizee")
+    assert attr_dict["ensemble_sizee"] == 7
 
 
 def test_reset(restrain_stage):
@@ -202,14 +202,20 @@ def setup_calc(mock_run_process):
                 ensemble_size=1,
                 stream_log_level=logging.CRITICAL,  # Silence the logging
             )
-            assert setup_calc.prep_stage == a3.run.enums.PreparationStage.PARAMETERISED
+            assert (
+                setup_calc.prep_stage.name
+                == a3.run.enums.PreparationStage.PARAMETERISED.name
+            )
             setup_calc.setup(slurm=False)
             yield setup_calc
 
     def test_setup_calc_overall(self, setup_calc, mock_run_process):
         """Test that setting up the calculation was successful at a high level."""
         assert setup_calc.setup_complete == True
-        assert setup_calc.prep_stage == a3.run.enums.PreparationStage.PREEQUILIBRATED
+        assert (
+            setup_calc.prep_stage.name
+            == a3.run.enums.PreparationStage.PREEQUILIBRATED.name
+        )
         assert len(setup_calc.legs) == 2
         legs = [leg.leg_type for leg in setup_calc.legs]
         assert a3.LegType.BOUND in legs
@@ -365,7 +371,9 @@ def test_integration_calculation(calc_slurm):
     """Integration test to check that all major stages of the calculation work."""
 
     # Check that the preparation stages work
-    assert calc_slurm.prep_stage == a3.run.enums.PreparationStage.PARAMETERISED
+    assert (
+        calc_slurm.prep_stage.name == a3.run.enums.PreparationStage.PARAMETERISED.name
+    )
     calc_slurm.setup(short_ensemble_equil=True)
     assert calc_slurm.setup_complete == True
     # Check that all required slurm bash jobs have been created