def test_can_get_nearby_atoms(self, mock_loc): model = Mock() mock_loc.return_value = (1, 2, 3) atom = Atom("C", 4, 8, 3) atom._model = model model.atoms_in_sphere.return_value = [1, 2, atom, 4] atoms = atom.nearby_atoms(4, a=1, b=2) model.atoms_in_sphere.assert_called_with(1, 2, 3, 4, a=1, b=2) self.assertEqual(atoms, [1, 2, 4])
def test_atom_with_no_model_has_no_nearby_atoms(self): atom = Atom("C", 4, 8, 3) self.assertEqual(atom.nearby_atoms(cutoff=1), set()) self.assertEqual(atom.nearby_atoms(cutoff=100), set())