Simplification of the logic for Legacy or atomistic Structure.

both in kpoints.py and test_kpoints.py

src/aiida/orm/nodes/data/array/kpoints.py

@@ -223,23 +223,18 @@ def set_cell_from_structure(self, structuredata):
 
         :param structuredata: an instance of StructureData
         """
-        from aiida.orm.nodes.data.structure import StructureData, has_atomistic
+        from aiida.orm.nodes.data.structure import has_atomistic, StructureData as LegacyStructureData
+
+        if has_atomistic():
+            structures_classes = (StructureData, LegacyStructureData)
+        else:
+            structures_classes = (LegacyStructureData,)
 
-        if not isinstance(structuredata, StructureData):
-            error_message = (
-                'An instance of StructureData or aiida-atomistic StructureData should be passed to '
-                'the KpointsData, found instead {}'.format(structuredata.__class__)
+        if not isinstance(structuredata, structures_classes):
+            raise TypeError(
+                f'An instance of {structures_classes} should be passed to '
+                f'the KpointsData, found instead {type(structuredata)}'
             )
-            if has_atomistic():
-                from aiida_atomistic import StructureData as AtomisticStructureData
-
-                if not isinstance(structuredata, AtomisticStructureData):
-                    raise ValueError(error_message)
-                else:
-                    cell = structuredata.cell
-                    self.set_cell(cell, structuredata.pbc)
-            else:
-                raise ValueError(error_message)
         else:
             cell = structuredata.cell
             self.set_cell(cell, structuredata.pbc)

tests/orm/nodes/data/test_kpoints.py

@@ -10,18 +10,22 @@
 
 import numpy as np
 import pytest
-from aiida.orm import KpointsData, StructureData, load_node
+from aiida.orm import KpointsData, load_node, StructureData as LegacyStructureData
 from aiida.orm.nodes.data.structure import has_atomistic
 
-skip_atomistic = pytest.mark.skipif(not has_atomistic(), reason='Unable to import aiida-atomistic')
-
-
+if not has_atomistic():
+    structures_classes = [LegacyStructureData,]
+else:
+    from aiida_atomistic import StructureData
+    structures_classes = [LegacyStructureData, StructureData]
+
+@pytest.mark.parametrize('structure_class', structures_classes)
 class TestKpoints:
     """Test for the `Kpointsdata` class."""
 
     @pytest.fixture(autouse=True)
-    def init_profile(self):
-        """Initialize the profile."""
+    def generate_structure(self, structure_class):
+        """Generate the StructureData."""
         alat = 5.430  # angstrom
         cell = [
             [
@@ -37,85 +41,13 @@ def init_profile(self):
             [0.5 * alat, 0.0, 0.5 * alat],
         ]
         self.alat = alat
-        structure = StructureData(cell=cell)
+        structure = LegacyStructureData(cell=cell)
         structure.append_atom(position=(0.000 * alat, 0.000 * alat, 0.000 * alat), symbols=['Si'])
         structure.append_atom(position=(0.250 * alat, 0.250 * alat, 0.250 * alat), symbols=['Si'])
-        self.structure = structure
-        # Define the expected reciprocal cell
-        val = 2.0 * np.pi / alat
-        self.expected_reciprocal_cell = np.array([[val, val, -val], [-val, val, val], [val, -val, val]])
-
-    def test_reciprocal_cell(self):
-        """Test the `reciprocal_cell` method.
-
-        This is a regression test for #2749.
-        """
-        kpt = KpointsData()
-        kpt.set_cell_from_structure(self.structure)
-
-        assert np.abs(kpt.reciprocal_cell - self.expected_reciprocal_cell).sum() == 0.0
-
-        # Check also after storing
-        kpt.store()
-        kpt2 = load_node(kpt.pk)
-        assert np.abs(kpt2.reciprocal_cell - self.expected_reciprocal_cell).sum() == 0.0
-
-    def test_get_kpoints(self):
-        """Test the `get_kpoints` method."""
-        kpt = KpointsData()
-        kpt.set_cell_from_structure(self.structure)
-
-        kpoints = [
-            [0.0, 0.0, 0.0],
-            [0.5, 0.5, 0.5],
-        ]
-
-        cartesian_kpoints = [
-            [0.0, 0.0, 0.0],
-            [np.pi / self.alat, np.pi / self.alat, np.pi / self.alat],
-        ]
-
-        kpt.set_kpoints(kpoints)
-        assert np.abs(kpt.get_kpoints() - np.array(kpoints)).sum() == 0.0
-        assert np.abs(kpt.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0
-
-        # Check also after storing
-        kpt.store()
-        kpt2 = load_node(kpt.pk)
-        assert np.abs(kpt2.get_kpoints() - np.array(kpoints)).sum() == 0.0
-        assert np.abs(kpt2.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0
-
-
-@skip_atomistic
-class TestKpointsAtomisticStructureData:
-    """Test for the `Kpointsdata` class using the new atomistic StructureData."""
-
-    @pytest.fixture(autouse=True)
-    def init_profile(self):
-        """Initialize the profile."""
-
-        from aiida_atomistic import StructureData, StructureDataMutable
-
-        alat = 5.430  # angstrom
-        cell = [
-            [
-                0.5 * alat,
-                0.5 * alat,
-                0.0,
-            ],
-            [
-                0.0,
-                0.5 * alat,
-                0.5 * alat,
-            ],
-            [0.5 * alat, 0.0, 0.5 * alat],
-        ]
-        self.alat = alat
-        mutable = StructureDataMutable()
-        mutable.set_cell(cell)
-        mutable.add_atom(positions=(0.000 * alat, 0.000 * alat, 0.000 * alat), symbols='Si')
-        mutable.add_atom(positions=(0.250 * alat, 0.250 * alat, 0.250 * alat), symbols='Si')
-        self.structure = StructureData.from_mutable(mutable)
+        if structure_class == LegacyStructureData:
+            self.structure = structure
+        else:
+            self.structure = LegacyStructureData.to_atomistic(structure)
         # Define the expected reciprocal cell
         val = 2.0 * np.pi / alat
         self.expected_reciprocal_cell = np.array([[val, val, -val], [-val, val, val], [val, -val, val]])
@@ -158,4 +90,4 @@ def test_get_kpoints(self):
         kpt.store()
         kpt2 = load_node(kpt.pk)
         assert np.abs(kpt2.get_kpoints() - np.array(kpoints)).sum() == 0.0
-        assert np.abs(kpt2.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0
+        assert np.abs(kpt2.get_kpoints(cartesian=True) - np.array(cartesian_kpoints)).sum() == 0.0