Add tests for the get_transfer_builder utility

tests/utils/conftest.py

@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*-
+"""A collection of useful pytest fixtures.
+
+* make_tempdir: returns a function to create temporary directories
+
+* makeget_computer: returns a function to create computers (note: these need to be destroyed after the test
+  by using 'clear_database' because as the python variables go out of scope, the tempdir used to create the
+  computers is deleted).
+
+"""
+import pytest
+
+# pylint: disable=redefined-outer-name,bad-option-value,raise-missing-from
+
+
+# Maybe this could go on aiida-core
+@pytest.fixture(scope='function')
+def make_tempdir(request):
+    """Provide a function to generate new temporary directories.
+
+    :return: The path to the directory
+    :rtype: str
+    """
+
+    def _make_tempdir(parent_dir=None):
+        import tempfile
+        import shutil
+
+        try:
+            dirpath = tempfile.mkdtemp(dir=parent_dir)
+        except FileNotFoundError as exc:
+            raise ValueError('The parent_dir provided was never created') from exc
+
+        # After the test function has completed, remove the directory again
+        # Design Note: yields would be a simpler way to do this, but can't use it
+        # inside a fixture factory without having problems with 'generator' object
+        def cleanup():
+            shutil.rmtree(dirpath)
+
+        request.addfinalizer(cleanup)
+
+        return dirpath
+
+    return _make_tempdir
+
+
+# Maybe this could go on aiida-core
+# It would be preferrable to do usefixtures but this just works for test, not other
+# fixtures (see https://github.com/pytest-dev/pytest/issues/3664).
+# We'll have to leave the pylint ignore until then
+#@pytest.mark.usefixtures('clear_database')
+# pylint: disable=unused-argument
+@pytest.fixture(scope='function')
+def makeget_computer(clear_database, make_tempdir):
+    """Provide a function to generate new computers.
+
+    :return: The computer node
+    :rtype: :py:class:`aiida.orm.Computer`
+    """
+
+    def _makeget_computer(label='localhost-test', transport_type='local'):
+        from aiida.orm import Computer
+        from aiida.common.exceptions import NotExistent
+
+        try:
+            computer = Computer.objects.get(label=label)
+
+        except NotExistent:
+            computer = Computer(
+                label=label,
+                description=f'{label} computer set up by test manager',
+                hostname=label,
+                workdir=make_tempdir(),
+                transport_type=transport_type,
+                scheduler_type='direct'
+            )
+            computer.store()
+            computer.set_minimum_job_poll_interval(0.)
+
+            if transport_type == 'local':
+                computer.configure()
+            else:
+                raise NotImplementedError('Transport `{transport_type}` not implemented')
+
+        return computer
+
+    return _makeget_computer

tests/utils/test_transfer.py

@@ -0,0 +1,140 @@
+# -*- coding: utf-8 -*-
+"""Unit tests for the :py:mod:`~aiida_quantumespresso.utils.transfer` module."""
+import pytest
+from packaging import version
+import aiida
+
+from aiida.orm import FolderData, RemoteData
+from aiida.engine import ProcessBuilder
+from aiida_quantumespresso.utils.transfer import get_transfer_builder
+
+pytestmark = pytest.mark.skipif(
+    version.parse(aiida.get_version()) < version.parse('1.6'),
+    reason='Transfer was released in AiiDA core v1.6',
+)
+
+
+@pytest.fixture(name='make_data_source')
+def fixture_make_data_source(make_tempdir):
+    """Provide a function to generate data sources specific for this test."""
+
+    def _make_data_source(computer=None, populate_level=0):
+        import os
+        import io
+
+        if computer is None:
+            node = FolderData()
+
+            filepaths = []
+            if populate_level > 0:
+                filepaths.append('data-file-schema.xml')
+                filepaths.append('charge-density.dat')
+            if populate_level > 1:
+                filepaths.append('paw.txt')
+
+            for filepath in filepaths:
+                node.put_object_from_filelike(io.StringIO(), path=filepath)
+
+        else:
+            remote_path = make_tempdir(parent_dir=computer.get_workdir())
+            node = RemoteData(computer=computer, remote_path=remote_path)
+
+            filepaths = []
+            if populate_level > 0:
+                filepaths.append('out/aiida.save/data-file-schema.xml')
+                filepaths.append('out/aiida.save/charge-density.dat')
+            if populate_level > 1:
+                filepaths.append('out/aiida.save/paw.txt')
+
+            for filepath in filepaths:
+                fullpath = os.path.join(remote_path, filepath)
+                os.makedirs(os.path.dirname(fullpath), exist_ok=True)
+                open(fullpath, 'w').close()
+
+        return node.store()
+
+    return _make_data_source
+
+
+@pytest.mark.parametrize('track', [True, False])
+@pytest.mark.parametrize('populate_level', [1, 2])
+@pytest.mark.parametrize('source_is_remote, provide_computer', [(True, True), (True, False), (False, True)])
+def test_transfer_builder_results(
+    makeget_computer, make_data_source, source_is_remote, provide_computer, populate_level, track
+):
+    """Test all viable input sets for the get_transfer_builder utility function."""
+
+    if provide_computer:
+        builder_computer = makeget_computer('computer1')
+        check_computer = builder_computer
+    else:
+        builder_computer = None
+
+    if source_is_remote:
+        folder_computer = makeget_computer('computer2')
+        data_source = make_data_source(computer=folder_computer, populate_level=populate_level)
+        retrieve_expected = True
+        expected_listname = 'symlink_files'
+        check_computer = folder_computer
+    else:
+        data_source = make_data_source(computer=None, populate_level=populate_level)
+        retrieve_expected = False
+        expected_listname = 'local_files'
+
+    if source_is_remote and provide_computer:
+        with pytest.warns(UserWarning) as warnings:
+            builder = get_transfer_builder(data_source, computer=builder_computer, track=track)
+        assert len(warnings) == 1
+        assert 'ignore' in str(warnings[0].message)
+        assert f'{builder_computer}' in str(warnings[0].message)
+        assert f'{folder_computer}' in str(warnings[0].message)
+        # Makes sure the information is there, but there is no generic way of checking the correctedness
+        # of the warning (i.e. which computer is the one selected) without checking for the exact text
+    else:
+        builder = get_transfer_builder(data_source, computer=builder_computer, track=track)
+
+    assert isinstance(builder, ProcessBuilder)
+    assert builder.metadata['computer'].pk == check_computer.pk
+    assert builder.source_nodes['source_node'] == data_source
+    assert builder.instructions.get_dict()['retrieve_files'] == retrieve_expected
+    assert expected_listname in builder.instructions.get_dict()
+
+    if source_is_remote:
+        # paw.txt are always loaded when remote copying
+        expected_setlist = {
+            ('source_node', 'out/aiida.save/data-file-schema.xml', 'data-file-schema.xml'),
+            ('source_node', 'out/aiida.save/charge-density.dat', 'charge-density.dat'),
+            ('source_node', 'out/aiida.save/paw.txt', 'paw.txt'),
+        }
+
+    else:
+        expected_setlist = {
+            ('source_node', 'data-file-schema.xml', 'out/aiida.save/data-file-schema.xml'),
+            ('source_node', 'charge-density.dat', 'out/aiida.save/charge-density.dat'),
+        }
+        if populate_level == 2:
+            expected_setlist.add(('source_node', 'paw.txt', 'out/aiida.save/paw.txt'))
+
+    # Note: I need to do the following transformation manually because sometimes what I get from:
+    #
+    # > builder.instructions.get_dict()[expected_listname]
+    #
+    # is a list of (unhashable) lists, and other times it returns a list of tupples
+    # This needs to be checked in aiida-core before changing the following to a simpler syntax
+    obtained_setlist = set()
+    for element in builder.instructions.get_dict()[expected_listname]:
+        obtained_setlist.add(tuple(element))
+    assert obtained_setlist == expected_setlist
+
+
+@pytest.mark.parametrize('track', [True, False])
+@pytest.mark.parametrize('populate_level', [1, 2])
+def test_transfer_builder_raise(make_data_source, populate_level, track):
+    """Test all raises for the get_transfer_builder utility function."""
+
+    data_source = make_data_source(computer=None, populate_level=populate_level)
+
+    with pytest.raises(ValueError) as execinfo:
+        _ = get_transfer_builder(data_source, computer=None, track=track)
+
+    assert 'computer' in str(execinfo.value)