fix unit test

materialsproject · Aug 5, 2024 · 2378f85 · 2378f85
1 parent 374f711
commit 2378f85
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/tests/io/lammps/test_data.py b/tests/io/lammps/test_data.py
@@ -512,11 +512,13 @@ def test_json_dict(self):
         pd.testing.assert_frame_equal(c2h6.atoms, self.ethane.atoms)
         ff = self.ethane.force_field
         rng = np.random.default_rng()
-        key, target_df = rng.choice(sorted(ff.items()), 1)[0]
+        ff_items = list(ff.items())
+        key, target_df = ff_items[rng.choice(len(ff_items))]
         c2h6.force_field[key].index = c2h6.force_field[key].index.map(int)
         assert pd.testing.assert_frame_equal(c2h6.force_field[key], target_df, check_dtype=False) is None, key
         topo = self.ethane.topology
-        key, target_df = rng.choice(sorted(topo.items()), 1)[0]
+        topo_items = list(topo.items())
+        key, target_df = topo_items[rng.choice(len(topo_items))]
         c2h6.topology[key].index = c2h6.topology[key].index.map(int)
         assert pd.testing.assert_frame_equal(c2h6.topology[key], target_df) is None, key
 
@@ -1066,24 +1068,28 @@ def test_json_dict(self):
         pd.testing.assert_frame_equal(lic3o3h4.atoms, self.li_ec.atoms)
         ff = self.li_ec.force_field
         rng = np.random.default_rng()
-        key, target_df = rng.choice(sorted(ff.items()), 1)[0]
+        ff_items = list(ff.items())
+        key, target_df = ff_items[rng.choice(len(ff_items))]
         lic3o3h4.force_field[key].index = lic3o3h4.force_field[key].index.map(int)
         assert pd.testing.assert_frame_equal(lic3o3h4.force_field[key], target_df, check_dtype=False) is None, key
         topo = self.li_ec.topology
-        key, target_df = rng.choice(sorted(topo.items()), 1)[0]
+        topo_items = list(topo.items())
+        key, target_df = topo_items[rng.choice(len(topo_items))]
         assert pd.testing.assert_frame_equal(lic3o3h4.topology[key], target_df) is None, key
         lic3o3h4.mols[1].masses.index = lic3o3h4.mols[1].masses.index.map(int)
         lic3o3h4.mols[1].atoms.index = lic3o3h4.mols[1].atoms.index.map(int)
         pd.testing.assert_frame_equal(lic3o3h4.mols[1].masses, self.li_ec.mols[1].masses)
         pd.testing.assert_frame_equal(lic3o3h4.mols[1].atoms, self.li_ec.mols[1].atoms)
         ff_1 = self.li_ec.mols[1].force_field
-        key, target_df = rng.choice(sorted(ff_1.items()), 1)[0]
+        ff1_items = list(ff_1.items())
+        key, target_df = ff1_items[rng.choice(len(ff1_items))]
         lic3o3h4.mols[1].force_field[key].index = lic3o3h4.mols[1].force_field[key].index.map(int)
         assert (
             pd.testing.assert_frame_equal(lic3o3h4.mols[1].force_field[key], target_df, check_dtype=False) is None
         ), key
         topo_1 = self.li_ec.mols[1].topology
-        key, target_df = rng.choice(sorted(topo_1.items()), 1)[0]
+        topo1_items = list(topo_1.items())
+        key, target_df = topo1_items[rng.choice(len(topo1_items))]
         lic3o3h4.mols[1].topology[key].index = lic3o3h4.mols[1].topology[key].index.map(int)
         assert pd.testing.assert_frame_equal(lic3o3h4.mols[1].topology[key], target_df) is None, key