def test_1(self): import os import shutil import tempfile from aiida.orm import DataFactory from aiida.orm import load_node from aiida.orm.calculation.job import JobCalculation from aiida.orm.importexport import export # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: StructureData = DataFactory('structure') sd = StructureData() sd.store() calc = JobCalculation() calc.set_computer(self.computer) calc.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc.store() calc.add_link_from(sd) pks = [sd.pk, calc.pk] attrs = {} for pk in pks: node = load_node(pk) attrs[node.uuid] = dict() for k in node.attrs(): attrs[node.uuid][k] = node.get_attr(k) filename = os.path.join(temp_folder, "export.tar.gz") export([calc.dbnode], outfile=filename, silent=True) self.clean_db() # NOTE: it is better to load new nodes by uuid, rather than assuming # that they will have the first 3 pks. In fact, a recommended policy in # databases is that pk always increment, even if you've deleted elements import_data(filename, silent=True) for uuid in attrs.keys(): node = load_node(uuid) # for k in node.attrs(): for k in attrs[uuid].keys(): self.assertEquals(attrs[uuid][k], node.get_attr(k)) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def test_complex_graph_import_export(self): """ This test checks that a small and bit complex graph can be correctly exported and imported. It will create the graph, store it to the database, export it to a file and import it. In the end it will check if the initial nodes are present at the imported graph. """ import tempfile import shutil import os from aiida.orm.calculation.job import JobCalculation from aiida.orm.data.folder import FolderData from aiida.orm.data.parameter import ParameterData from aiida.orm.data.remote import RemoteData from aiida.common.links import LinkType from aiida.orm.importexport import export, import_data from aiida.orm.utils import load_node from aiida.common.exceptions import NotExistent temp_folder = tempfile.mkdtemp() try: calc1 = JobCalculation() calc1.set_computer(self.computer) calc1.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc1.label = "calc1" calc1.store() calc1._set_state(u'RETRIEVING') pd1 = ParameterData() pd1.label = "pd1" pd1.store() pd2 = ParameterData() pd2.label = "pd2" pd2.store() rd1 = RemoteData() rd1.label = "rd1" rd1.set_remote_path("/x/y.py") rd1.set_computer(self.computer) rd1.store() rd1.add_link_from(calc1, link_type=LinkType.CREATE) calc2 = JobCalculation() calc2.set_computer(self.computer) calc2.set_resources({ "num_machines": 1, "num_mpiprocs_per_machine": 1 }) calc2.label = "calc2" calc2.store() calc2.add_link_from(pd1, link_type=LinkType.INPUT) calc2.add_link_from(pd2, link_type=LinkType.INPUT) calc2.add_link_from(rd1, link_type=LinkType.INPUT) calc2._set_state(u'SUBMITTING') fd1 = FolderData() fd1.label = "fd1" fd1.store() fd1.add_link_from(calc2, link_type=LinkType.CREATE) node_uuids_labels = { calc1.uuid: calc1.label, pd1.uuid: pd1.label, pd2.uuid: pd2.label, rd1.uuid: rd1.label, calc2.uuid: calc2.label, fd1.uuid: fd1.label } filename = os.path.join(temp_folder, "export.tar.gz") export([fd1.dbnode], outfile=filename, silent=True) self.clean_db() import_data(filename, silent=True, ignore_unknown_nodes=True) for uuid, label in node_uuids_labels.iteritems(): try: load_node(uuid) except NotExistent: self.fail("Node with UUID {} and label {} was not " "found.".format(uuid, label)) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def test_6(self): """ This test checks that nodes belonging to user A (which is not the default user) can be correctly exported, imported, enriched with nodes from the default user, re-exported & re-imported and that in the end all the nodes that have been finally imported belonging to the right users. """ import os import shutil import tempfile from aiida.orm import load_node from aiida.orm.calculation.job import JobCalculation from aiida.orm.data.structure import StructureData from aiida.orm.importexport import export from aiida.common.datastructures import calc_states from aiida.common.links import LinkType from aiida.common.utils import get_configured_user_email from aiida.orm.user import User # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: # Create another user new_email = "[email protected]" user = User(email=new_email) user.force_save() # Create a structure data node that has a calculation as output sd1 = StructureData() sd1.dbnode.user = user._dbuser sd1.label = 'sd1' sd1.store() jc1 = JobCalculation() jc1.set_computer(self.computer) jc1.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc1.dbnode.user = user._dbuser jc1.label = 'jc1' jc1.store() jc1.add_link_from(sd1) jc1._set_state(calc_states.PARSING) # Create some nodes from a different user sd2 = StructureData() sd2.dbnode.user = user._dbuser sd2.label = 'sd2' sd2.store() sd2.add_link_from(jc1, label='l1', link_type=LinkType.RETURN) # Set the jc1 to FINISHED jc1._set_state(calc_states.FINISHED) # At this point we export the generated data filename1 = os.path.join(temp_folder, "export1.tar.gz") export([sd2.dbnode], outfile=filename1, silent=True) uuids1 = [sd1.uuid, jc1.uuid, sd2.uuid] self.clean_db() self.insert_data() import_data(filename1, silent=True) # Check that the imported nodes are correctly imported and that # the user assigned to the nodes is the right one for uuid in uuids1: self.assertEquals(load_node(uuid).get_user().email, new_email) # Now we continue to generate more data based on the imported # data sd2_imp = load_node(sd2.uuid) jc2 = JobCalculation() jc2.set_computer(self.computer) jc2.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc2.label = 'jc2' jc2.store() jc2.add_link_from(sd2_imp, label='l2') jc2._set_state(calc_states.PARSING) sd3 = StructureData() sd3.label = 'sd3' sd3.store() sd3.add_link_from(jc2, label='l3', link_type=LinkType.RETURN) # Set the jc2 to FINISHED jc2._set_state(calc_states.FINISHED) # Store the UUIDs of the nodes that should be checked # if they can be imported correctly. uuids2 = [jc2.uuid, sd3.uuid] filename2 = os.path.join(temp_folder, "export2.tar.gz") export([sd3.dbnode], outfile=filename2, silent=True) self.clean_db() self.insert_data() import_data(filename2, silent=True) # Check that the imported nodes are correctly imported and that # the user assigned to the nodes is the right one for uuid in uuids1: self.assertEquals(load_node(uuid).get_user().email, new_email) for uuid in uuids2: self.assertEquals(load_node(uuid).get_user().email, get_configured_user_email()) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)
def test_5(self): """ This test checks that nodes belonging to different users are correctly exported & imported. """ import os import shutil import tempfile from aiida.orm import load_node from aiida.orm.calculation.job import JobCalculation from aiida.orm.data.structure import StructureData from aiida.orm.importexport import export from aiida.common.datastructures import calc_states from aiida.common.links import LinkType from aiida.orm.user import User from aiida.common.utils import get_configured_user_email # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() try: # Create another user new_email = "[email protected]" user = User(email=new_email) user.force_save() # Create a structure data node that has a calculation as output sd1 = StructureData() sd1.dbnode.user = user._dbuser sd1.label = 'sd1' sd1.store() jc1 = JobCalculation() jc1.set_computer(self.computer) jc1.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc1.dbnode.user = user._dbuser jc1.label = 'jc1' jc1.store() jc1.add_link_from(sd1) jc1._set_state(calc_states.PARSING) # Create some nodes from a different user sd2 = StructureData() sd2.dbnode.user = user._dbuser sd2.label = 'sd2' sd2.store() sd2.add_link_from(jc1, label='l1', link_type=LinkType.RETURN) jc2 = JobCalculation() jc2.set_computer(self.computer) jc2.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) jc2.label = 'jc2' jc2.store() jc2.add_link_from(sd2, label='l2') jc2._set_state(calc_states.PARSING) sd3 = StructureData() sd3.label = 'sd3' sd3.store() sd3.add_link_from(jc2, label='l3', link_type=LinkType.RETURN) uuids_u1 = [sd1.uuid, jc1.uuid, sd2.uuid] uuids_u2 = [jc2.uuid, sd3.uuid] filename = os.path.join(temp_folder, "export.tar.gz") export([sd3.dbnode], outfile=filename, silent=True) self.clean_db() import_data(filename, silent=True) # Check that the imported nodes are correctly imported and that # the user assigned to the nodes is the right one for uuid in uuids_u1: self.assertEquals(load_node(uuid).get_user().email, new_email) for uuid in uuids_u2: self.assertEquals(load_node(uuid).get_user().email, get_configured_user_email()) finally: # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True)