Update test_cli and test_analysis

SMTG-Bham · Feb 6, 2025 · 90202f3 · 90202f3
1 parent bdf35e5
commit 90202f3
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 64 deletions.
diff --git a/shakenbreak/input.py b/shakenbreak/input.py
@@ -2115,16 +2115,16 @@ def _update_distortion_metadata(
         distortion_metadata: dict,
         defect_name: str,
         charge: int,
-        defect_site_index: int,
         num_nearest_neighbours: int,
         distorted_atoms: list,
+        defect_site_index: Optional[int] = None,
         defect_type: str = "",
     ) -> dict:
         """
         Update distortion_metadata with distortion information for each
         charged defect.
         """
-        if defect_site_index:  # store site index of defect if not vacancy
+        if defect_site_index is not None:  # store site index of defect if not vacancy
             distortion_metadata["defects"][defect_name]["defect_site_index"] = defect_site_index
         rattle_parameters = self._mc_rattle_kwargs.copy()
         rattle_parameters["stdev"] = self.stdev
@@ -2429,10 +2429,10 @@ def apply_distortions(
                     distortion_metadata=self.distortion_metadata,
                     defect_name=defect_name,
                     charge=charge,
-                    defect_site_index=defect_site_index,
                     num_nearest_neighbours=num_nearest_neighbours,
                     distorted_atoms=distorted_atoms,
                     defect_type=defect_type,
+                    defect_site_index=defect_site_index,
                 )
 
         return distorted_defects_dict, self.distortion_metadata

diff --git a/tests/test_analysis.py b/tests/test_analysis.py
@@ -289,7 +289,7 @@ def test_analyse_defect_site(self):
             output = analysis.analyse_defect_site(
                 self.Int_Cd_2_minus0pt6_NN_10_struc_rattled,
                 name="Int_Cd_2",
-                site_num=65,
+                site_index=64,
             )
             # mock_print.assert_any_call("==> ", "Int_Cd_2 structural analysis ", " <==")
             mock_print.assert_any_call(
@@ -340,7 +340,7 @@ def test_analyse_defect_site(self):
             analysis.analyse_defect_site(
                 self.Int_Cd_2_minus0pt6_NN_10_struc_rattled,
                 name="Int_Cd_2",
-                # no site_num or vac_coords specified
+                # no site_index or vac_coords specified
             )
 
     def test_get_structures(self):