test(unit): add tests for Neighbors.get_symmetry

Author: harryswift01

tests/unit/CodeEntropy/levels/test_neighbors.py

@@ -1,5 +1,5 @@
 import contextlib
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, call, patch
 
 import numpy as np
 import pytest
@@ -235,3 +235,125 @@ class HybridizationType:
         result = neighbors._get_linear(rdkit_mol, number_heavy)
 
     assert result is False
+
+
+def test_get_symmetry_returns_dicts_for_single_group():
+    neighbors = Neighbors()
+    universe = MagicMock()
+    groups = {7: [42, 99]}
+
+    rdkit_mol = MagicMock()
+
+    neighbors._get_rdkit_mol = MagicMock(return_value=(rdkit_mol, 5, 8))
+    neighbors._get_symmetry_number = MagicMock(return_value=12)
+    neighbors._get_linear = MagicMock(return_value=True)
+
+    symmetry_number, linear = neighbors.get_symmetry(universe, groups)
+
+    assert symmetry_number == {7: 12}
+    assert linear == {7: True}
+
+    neighbors._get_rdkit_mol.assert_called_once_with(universe, 42)
+    neighbors._get_symmetry_number.assert_called_once_with(rdkit_mol, 5, 8)
+    neighbors._get_linear.assert_called_once_with(rdkit_mol, 5)
+
+
+def test_get_symmetry_uses_first_molecule_in_each_group_only():
+    neighbors = Neighbors()
+    universe = MagicMock()
+    groups = {
+        0: [10, 11, 12],
+        1: [20, 21],
+    }
+
+    rdkit_mol_0 = MagicMock()
+    rdkit_mol_1 = MagicMock()
+
+    neighbors._get_rdkit_mol = MagicMock(
+        side_effect=[
+            (rdkit_mol_0, 3, 6),
+            (rdkit_mol_1, 4, 8),
+        ]
+    )
+    neighbors._get_symmetry_number = MagicMock(side_effect=[2, 4])
+    neighbors._get_linear = MagicMock(side_effect=[False, True])
+
+    symmetry_number, linear = neighbors.get_symmetry(universe, groups)
+
+    assert symmetry_number == {0: 2, 1: 4}
+    assert linear == {0: False, 1: True}
+
+    assert neighbors._get_rdkit_mol.call_args_list == [
+        call(universe, 10),
+        call(universe, 20),
+    ]
+
+
+def test_get_symmetry_calls_helpers_for_each_group_in_order():
+    neighbors = Neighbors()
+    universe = MagicMock()
+    groups = {
+        3: [100],
+        5: [200],
+    }
+
+    rdkit_mol_a = MagicMock()
+    rdkit_mol_b = MagicMock()
+
+    neighbors._get_rdkit_mol = MagicMock(
+        side_effect=[
+            (rdkit_mol_a, 1, 2),
+            (rdkit_mol_b, 7, 0),
+        ]
+    )
+    neighbors._get_symmetry_number = MagicMock(side_effect=[9, 1])
+    neighbors._get_linear = MagicMock(side_effect=[True, False])
+
+    symmetry_number, linear = neighbors.get_symmetry(universe, groups)
+
+    assert symmetry_number == {3: 9, 5: 1}
+    assert linear == {3: True, 5: False}
+
+    assert neighbors._get_symmetry_number.call_args_list == [
+        call(rdkit_mol_a, 1, 2),
+        call(rdkit_mol_b, 7, 0),
+    ]
+    assert neighbors._get_linear.call_args_list == [
+        call(rdkit_mol_a, 1),
+        call(rdkit_mol_b, 7),
+    ]
+
+
+def test_get_symmetry_returns_empty_dicts_for_empty_groups():
+    neighbors = Neighbors()
+    universe = MagicMock()
+    groups = {}
+
+    neighbors._get_rdkit_mol = MagicMock()
+    neighbors._get_symmetry_number = MagicMock()
+    neighbors._get_linear = MagicMock()
+
+    symmetry_number, linear = neighbors.get_symmetry(universe, groups)
+
+    assert symmetry_number == {}
+    assert linear == {}
+
+    neighbors._get_rdkit_mol.assert_not_called()
+    neighbors._get_symmetry_number.assert_not_called()
+    neighbors._get_linear.assert_not_called()
+
+
+def test_get_symmetry_propagates_error_from_get_rdkit_mol():
+    neighbors = Neighbors()
+    universe = MagicMock()
+    groups = {0: [123]}
+
+    neighbors._get_rdkit_mol = MagicMock(side_effect=RuntimeError("bad molecule"))
+    neighbors._get_symmetry_number = MagicMock()
+    neighbors._get_linear = MagicMock()
+
+    with pytest.raises(RuntimeError, match="bad molecule"):
+        neighbors.get_symmetry(universe, groups)
+
+    neighbors._get_symmetry_number.assert_not_called()
+    neighbors._get_linear.assert_not_called()