def test_transaction(self):
"""Test that transaction nesting works."""
user1 = orm.User('*****@*****.**').store()
user2 = orm.User('*****@*****.**').store()
try:
with self.backend.transaction():
user1.email = '*****@*****.**'
user2.email = '*****@*****.**'
raise RuntimeError
except RuntimeError:
pass
self.assertEqual(user1.email, '*****@*****.**')
self.assertEqual(user2.email, '*****@*****.**')
def test_group_import_existing(self, temp_dir):
"""
Testing what happens when I try to import a group that already exists in the
database. This should raise an appropriate exception
"""
grouplabel = 'node_group_existing'
# Create another user
new_email = '[email protected]'
user = orm.User(email=new_email)
user.store()
# Create a structure data node
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd'
sd1.store()
# Create a group and add the data inside
group = orm.Group(label=grouplabel)
group.store()
group.add_nodes([sd1])
# At this point we export the generated data
filename = os.path.join(temp_dir, 'export1.tar.gz')
export([group], outfile=filename, silent=True)
self.clean_db()
self.insert_data()
# Creating a group of the same name
group = orm.Group(label='node_group_existing')
group.store()
import_data(filename, silent=True)
# The import should have created a new group with a suffix
# I check for this:
builder = orm.QueryBuilder().append(
orm.Group, filters={'label': {
'like': grouplabel + '%'
}})
self.assertEqual(builder.count(), 2)
# Now I check for the group having one member, and whether the name is different:
builder = orm.QueryBuilder()
builder.append(orm.Group,
filters={'label': {
'like': grouplabel + '%'
}},
tag='g',
project='label')
builder.append(orm.StructureData, with_group='g')
self.assertEqual(builder.count(), 1)
# I check that the group name was changed:
self.assertTrue(builder.all()[0][0] != grouplabel)
# I import another name, the group should not be imported again
import_data(filename, silent=True)
builder = orm.QueryBuilder()
builder.append(orm.Group,
filters={'label': {
'like': grouplabel + '%'
}})
self.assertEqual(builder.count(), 2)
def test_get_upf_groups(self):
"""Test the `UpfData.get_upf_groups` class method."""
type_string = orm.GroupTypeString.UPFGROUP_TYPE.value
label_01 = 'family_01'
label_02 = 'family_02'
user = orm.User(email='alternate@localhost').store()
self.assertEqual(orm.UpfData.get_upf_groups(), [])
# Create group with default user and add `Ba` pseudo
family_01, _ = orm.Group.objects.get_or_create(label=label_01,
type_string=type_string)
family_01.add_nodes([self.pseudo_barium])
family_01.store()
self.assertEqual(
{group.label
for group in orm.UpfData.get_upf_groups()}, {label_01})
# Create group with different user and add `O` pseudo
family_02, _ = orm.Group.objects.get_or_create(label=label_02,
type_string=type_string,
user=user)
family_02.add_nodes([self.pseudo_oxygen])
family_02.store()
self.assertEqual(
{group.label
for group in orm.UpfData.get_upf_groups()}, {label_01, label_02})
# Filter on a given user
self.assertEqual(
{
group.label
for group in orm.UpfData.get_upf_groups(user=user.email)
}, {label_02})
# Filter on a given element
groups = {
group.label
for group in orm.UpfData.get_upf_groups(filter_elements='O')
}
self.assertEqual(groups, {label_02})
# Filter on a given element and user
groups = {
group.label
for group in orm.UpfData.get_upf_groups(filter_elements='O',
user=user.email)
}
self.assertEqual(groups, {label_02})
# Filter on element and user that should not match anything
groups = {
group.label
for group in orm.UpfData.get_upf_groups(filter_elements='Ba',
user=user.email)
}
self.assertEqual(groups, set([]))
def test_store_in_transaction(self):
"""Test that storing inside a transaction is correctly dealt with."""
user1 = orm.User('*****@*****.**')
with self.backend.transaction():
user1.store()
# the following shouldn't raise
orm.User.objects.get(email='*****@*****.**')
user2 = orm.User('*****@*****.**')
try:
with self.backend.transaction():
user2.store()
raise RuntimeError
except RuntimeError:
pass
with self.assertRaises(exceptions.NotExistent):
orm.User.objects.get(email='*****@*****.**')
def setUp(self):
super().setUp()
created, user = orm.User.objects.get_or_create(email=USER_1['email'])
for key, _ in USER_1.items():
if key != 'email':
setattr(user, key, USER_1[key])
if created:
orm.User(**USER_1).store()
self.cli_runner = CliRunner()
def setUp(self):
super(TestVerdiUserCommand, self).setUp()
created, user = orm.User.objects.get_or_create(email=user_1['email'])
for key, value in user_1.items():
if key != 'email':
setattr(user, key, user_1[key])
if created:
orm.User(**user_1).store()
self.cli_runner = CliRunner()
def create_user(self):
"""This method creates and stores the default user. It has the same effect
as the verdi setup."""
from aiida.manage.configuration import get_config
self.user_email = get_config().current_profile.default_user
# Since the default user is needed for many operations in AiiDA, it is not deleted by clean_db.
# In principle, it should therefore always exist - if not we create it anyhow.
try:
self.user = orm.User.objects.get(email=self.user_email)
except exceptions.NotExistent:
self.user = orm.User(email=self.user_email).store()
def test_nodes_in_group(self, temp_dir):
"""
This test checks that nodes that belong to a specific group are
correctly imported and exported.
"""
from aiida.common.links import LinkType
# Create another user
new_email = '[email protected]'
user = orm.User(email=new_email)
user.store()
# Create a structure data node that has a calculation as output
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd1'
sd1.store()
jc1 = orm.CalcJobNode()
jc1.computer = self.computer
jc1.set_option('resources', {
'num_machines': 1,
'num_mpiprocs_per_machine': 1
})
jc1.user = user
jc1.label = 'jc1'
jc1.add_incoming(sd1, link_type=LinkType.INPUT_CALC, link_label='link')
jc1.store()
jc1.seal()
# Create a group and add the data inside
gr1 = orm.Group(label='node_group')
gr1.store()
gr1.add_nodes([sd1, jc1])
gr1_uuid = gr1.uuid
# At this point we export the generated data
filename1 = os.path.join(temp_dir, 'export1.tar.gz')
export([sd1, jc1, gr1], outfile=filename1, silent=True)
n_uuids = [sd1.uuid, jc1.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 n_uuids:
self.assertEqual(orm.load_node(uuid).user.email, new_email)
# Check that the exported group is imported correctly
builder = orm.QueryBuilder()
builder.append(orm.Group, filters={'uuid': {'==': gr1_uuid}})
self.assertEqual(builder.count(), 1, 'The group was not found.')
def test_transaction_nesting(self):
"""Test that transaction nesting works."""
user = orm.User('*****@*****.**').store()
with self.backend.transaction():
user.email = '*****@*****.**'
try:
with self.backend.transaction():
user.email = '*****@*****.**'
self.assertEqual(user.email, '*****@*****.**')
raise RuntimeError
except RuntimeError:
pass
self.assertEqual(user.email, '*****@*****.**')
self.assertEqual(user.email, '*****@*****.**')
def convert(self, value, param, ctx):
from aiida import orm
results = orm.User.objects.find({'email': value})
if not results:
if self._create:
return orm.User(email=value)
self.fail(f"User '{value}' not found", param, ctx)
if len(results) > 1:
self.fail(f"Multiple users found with email '{value}': {results}")
return results[0]
def test_exclude_comments_flag(self, temp_dir):
"""Test comments and associated commenting users are not exported when using `include_comments=False`."""
# Create users, node, and comments
user_one = orm.User.objects.get_default()
user_two = orm.User(email='[email protected]').store()
node = orm.Data().store()
orm.Comment(node, user_one, self.comments[0]).store()
orm.Comment(node, user_one, self.comments[1]).store()
orm.Comment(node, user_two, self.comments[2]).store()
orm.Comment(node, user_two, self.comments[3]).store()
# Get values prior to export
users_email = [u.email for u in [user_one, user_two]]
node_uuid = node.uuid
# Check that node belongs to user_one
self.assertEqual(node.user.email, users_email[0])
# Export nodes, excluding comments
export_file = os.path.join(temp_dir, 'export.tar.gz')
export([node],
outfile=export_file,
silent=True,
include_comments=False)
# Clean database and reimport exported file
self.reset_database()
import_data(export_file, silent=True)
# Get node, users, and comments
import_nodes = orm.QueryBuilder().append(orm.Node,
project=['uuid']).all()
import_comments = orm.QueryBuilder().append(orm.Comment,
project=['uuid']).all()
import_users = orm.QueryBuilder().append(orm.User,
project=['email']).all()
# There should be exactly: 1 Node, 0 Comments, 1 User
self.assertEqual(len(import_nodes), 1)
self.assertEqual(len(import_comments), 0)
self.assertEqual(len(import_users), 1)
# Check it's the correct user (and node)
self.assertEqual(str(import_nodes[0][0]), node_uuid)
self.assertEqual(str(import_users[0][0]), users_email[0])
def test_group_export(self, temp_dir):
"""Exporting a group includes its extras and nodes."""
# Create a new user
new_email = '[email protected]'
user = orm.User(email=new_email)
user.store()
# Create a structure data node
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd1'
sd1.store()
# Create a group and add the node
group = orm.Group(label='node_group')
group.set_extra('test', 1)
group.store()
group.add_nodes([sd1])
group_uuid = group.uuid
# Export the generated data, clean the database and import it again
filename = os.path.join(temp_dir, 'export.aiida')
export([group], filename=filename, silent=True)
n_uuids = [sd1.uuid]
self.clean_db()
self.insert_data()
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 n_uuids:
self.assertEqual(orm.load_node(uuid).user.email, new_email)
# Check that the exported group is imported correctly
builder = orm.QueryBuilder()
builder.append(orm.Group, filters={'uuid': {'==': group_uuid}})
self.assertEqual(builder.count(), 1, 'The group was not found.')
imported_group = builder.all()[0][0]
self.assertEqual(imported_group.get_extra('test'), 1, 'Extra missing on imported group')
def test_group_export(self, temp_dir):
"""Test that when exporting just a group, its nodes are also exported"""
# Create another user
new_email = '[email protected]'
user = orm.User(email=new_email)
user.store()
# Create a structure data node
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd1'
sd1.store()
# Create a group and add the data inside
group = orm.Group(label='node_group')
group.store()
group.add_nodes([sd1])
group_uuid = group.uuid
# At this point we export the generated data
filename = os.path.join(temp_dir, 'export.tar.gz')
export([group], outfile=filename, silent=True)
n_uuids = [sd1.uuid]
self.clean_db()
self.insert_data()
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 n_uuids:
self.assertEqual(orm.load_node(uuid).user.email, new_email)
# Check that the exported group is imported correctly
builder = orm.QueryBuilder()
builder.append(orm.Group, filters={'uuid': {'==': group_uuid}})
self.assertEqual(builder.count(), 1, 'The group was not found.')
def test_comment_querybuilder(self):
# pylint: disable=too-many-locals
"""Test querying for comments by joining on nodes in the QueryBuilder."""
user_one = self.user
user_two = orm.User(email='[email protected]').store()
node_one = orm.Data().store()
comment_one = Comment(node_one, user_one, 'comment_one').store()
node_two = orm.Data().store()
comment_two = Comment(node_two, user_one, 'comment_two').store()
comment_three = Comment(node_two, user_one, 'comment_three').store()
node_three = orm.CalculationNode().store()
comment_four = Comment(node_three, user_two,
'new_user_comment').store()
node_four = orm.CalculationNode().store()
comment_five = Comment(node_four, user_one, 'user one comment').store()
comment_six = Comment(node_four, user_two, 'user two comment').store()
# Retrieve a node by joining on a specific comment
builder = orm.QueryBuilder()
builder.append(Comment, tag='comment', filters={'id': comment_one.id})
builder.append(orm.Node, with_comment='comment', project=['uuid'])
nodes = builder.all()
self.assertEqual(len(nodes), 1)
for node in nodes:
self.assertIn(str(node[0]), [node_one.uuid])
# Retrieve a comment by joining on a specific node
builder = orm.QueryBuilder()
builder.append(orm.Node, tag='node', filters={'id': node_two.id})
builder.append(Comment, with_node='node', project=['uuid'])
comments = builder.all()
self.assertEqual(len(comments), 2)
for comment in comments:
self.assertIn(str(comment[0]),
[comment_two.uuid, comment_three.uuid])
# Retrieve a user by joining on a specific comment
builder = orm.QueryBuilder()
builder.append(Comment, tag='comment', filters={'id': comment_four.id})
builder.append(orm.User, with_comment='comment', project=['email'])
users = builder.all()
self.assertEqual(len(users), 1)
for user in users:
self.assertEqual(str(user[0]), user_two.email)
# Retrieve a comment by joining on a specific user
builder = orm.QueryBuilder()
builder.append(orm.User, tag='user', filters={'email': user_one.email})
builder.append(Comment, with_user='******', project=['uuid'])
comments = builder.all()
self.assertEqual(len(comments), 5)
for comment in comments:
self.assertIn(str(comment[0]), [
self.comment.uuid, comment_one.uuid, comment_two.uuid,
comment_three.uuid, comment_five.uuid
])
# Retrieve users from comments of a single node by joining specific node
builder = orm.QueryBuilder()
builder.append(orm.Node, tag='node', filters={'id': node_four.id})
builder.append(Comment,
tag='comments',
with_node='node',
project=['uuid'])
builder.append(orm.User, with_comment='comments', project=['email'])
comments_and_users = builder.all()
self.assertEqual(len(comments_and_users), 2)
for entry in comments_and_users:
self.assertEqual(len(entry), 2)
comment_uuid = str(entry[0])
user_email = str(entry[1])
self.assertIn(comment_uuid, [comment_five.uuid, comment_six.uuid])
self.assertIn(user_email, [user_one.email, user_two.email])
def test_nodes_belonging_to_different_users(self, temp_dir):
"""
This test checks that nodes belonging to different users are correctly
exported & imported.
"""
from aiida.common.links import LinkType
from aiida.manage.manager import get_manager
manager = get_manager()
# Create another user
new_email = '[email protected]'
user = orm.User(email=new_email).store()
# Create a structure data node that has a calculation as output
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd1'
sd1.store()
jc1 = orm.CalcJobNode()
jc1.computer = self.computer
jc1.set_option('resources', {
'num_machines': 1,
'num_mpiprocs_per_machine': 1
})
jc1.user = user
jc1.label = 'jc1'
jc1.add_incoming(sd1, link_type=LinkType.INPUT_CALC, link_label='link')
jc1.store()
# Create some nodes from a different user
sd2 = orm.StructureData()
sd2.user = user
sd2.label = 'sd2'
sd2.store()
sd2.add_incoming(jc1, link_type=LinkType.CREATE,
link_label='l1') # I assume jc1 CREATED sd2
jc1.seal()
jc2 = orm.CalcJobNode()
jc2.computer = self.computer
jc2.set_option('resources', {
'num_machines': 1,
'num_mpiprocs_per_machine': 1
})
jc2.label = 'jc2'
jc2.add_incoming(sd2, link_type=LinkType.INPUT_CALC, link_label='l2')
jc2.store()
sd3 = orm.StructureData()
sd3.label = 'sd3'
sd3.store()
sd3.add_incoming(jc2, link_type=LinkType.CREATE, link_label='l3')
jc2.seal()
uuids_u1 = [sd1.uuid, jc1.uuid, sd2.uuid]
uuids_u2 = [jc2.uuid, sd3.uuid]
filename = os.path.join(temp_dir, 'export.aiida')
export([sd3], filename=filename, silent=True)
self.clean_db()
self.create_user()
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:
node = orm.load_node(uuid=uuid)
self.assertEqual(node.user.email, new_email)
for uuid in uuids_u2:
self.assertEqual(
orm.load_node(uuid).user.email,
manager.get_profile().default_user)
def test_non_default_user_nodes(self, temp_dir): # pylint: disable=too-many-statements
"""
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.
"""
from aiida.common.links import LinkType
from aiida.manage.manager import get_manager
manager = get_manager()
# Create another user
new_email = '[email protected]'
user = orm.User(email=new_email).store()
# Create a structure data node that has a calculation as output
sd1 = orm.StructureData()
sd1.user = user
sd1.label = 'sd1'
sd1.store()
jc1 = orm.CalcJobNode()
jc1.computer = self.computer
jc1.set_option('resources', {
'num_machines': 1,
'num_mpiprocs_per_machine': 1
})
jc1.user = user
jc1.label = 'jc1'
jc1.add_incoming(sd1, link_type=LinkType.INPUT_CALC, link_label='link')
jc1.store()
# Create some nodes from a different user
sd2 = orm.StructureData()
sd2.user = user
sd2.label = 'sd2'
sd2.add_incoming(jc1, link_type=LinkType.CREATE, link_label='l1')
sd2.store()
jc1.seal()
sd2_uuid = sd2.uuid
# At this point we export the generated data
filename1 = os.path.join(temp_dir, 'export1.aiidaz')
export([sd2], filename=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.assertEqual(orm.load_node(uuid).user.email, new_email)
# Now we continue to generate more data based on the imported
# data
sd2_imp = orm.load_node(sd2_uuid)
jc2 = orm.CalcJobNode()
jc2.computer = self.computer
jc2.set_option('resources', {
'num_machines': 1,
'num_mpiprocs_per_machine': 1
})
jc2.label = 'jc2'
jc2.add_incoming(sd2_imp,
link_type=LinkType.INPUT_CALC,
link_label='l2')
jc2.store()
sd3 = orm.StructureData()
sd3.label = 'sd3'
sd3.add_incoming(jc2, link_type=LinkType.CREATE, link_label='l3')
sd3.store()
jc2.seal()
# 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_dir, 'export2.aiida')
export([sd3], filename=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.assertEqual(orm.load_node(uuid).user.email, new_email)
for uuid in uuids2:
self.assertEqual(
orm.load_node(uuid).user.email,
manager.get_profile().default_user)
def test_multiple_user_comments_single_node(self, temp_dir):
""" Test multiple users commenting on a single orm.CalculationNode """
# Create users, node, and comments
user_one = orm.User.objects.get_default()
user_two = orm.User(email='[email protected]').store()
node = orm.CalculationNode().store()
node.seal()
comment_one = orm.Comment(node, user_one, self.comments[0]).store()
comment_two = orm.Comment(node, user_one, self.comments[1]).store()
comment_three = orm.Comment(node, user_two, self.comments[2]).store()
comment_four = orm.Comment(node, user_two, self.comments[3]).store()
# Get values prior to export
users_email = [u.email for u in [user_one, user_two]]
node_uuid = str(node.uuid)
user_one_comments_uuid = [
str(c.uuid) for c in [comment_one, comment_two]
]
user_two_comments_uuid = [
str(c.uuid) for c in [comment_three, comment_four]
]
# Export node, along with comments and users recursively
export_file = os.path.join(temp_dir, 'export.tar.gz')
export([node], outfile=export_file, silent=True)
# Clean database and reimport exported file
self.reset_database()
import_data(export_file, silent=True)
# Get node, users, and comments
builder = orm.QueryBuilder()
builder.append(orm.Node, tag='node', project=['uuid'])
builder.append(orm.Comment,
tag='comment',
with_node='node',
project=['uuid'])
builder.append(orm.User, with_comment='comment', project=['email'])
entries = builder.all()
# Check that all 4 comments are retrieved, along with their respective node and user
self.assertEqual(len(entries), len(self.comments))
# Go through [Node.uuid, Comment.uuid, User.email]-entries
imported_node_uuids = set()
imported_user_one_comment_uuids = set()
imported_user_two_comment_uuids = set()
imported_user_emails = set()
for entry in entries:
self.assertEqual(len(entry), 3) # 1 Node + 1 Comment + 1 User
# Add node to set of imported nodes
imported_node_uuids.add(str(entry[0]))
# Add user to set of imported users
import_user_email = entry[2]
imported_user_emails.add(str(import_user_email))
# Add comment to set of imported comments pertaining to correct user
if import_user_email == users_email[0]:
# User_one comments
imported_user_one_comment_uuids.add(str(entry[1]))
else:
# User_two comments
imported_user_two_comment_uuids.add(str(entry[1]))
# Check same number of nodes (1) and users (2) were ex- and imported
self.assertEqual(len(imported_node_uuids), 1)
self.assertEqual(len(imported_user_emails), len(users_email))
# Check imported node equals exported node
self.assertSetEqual(imported_node_uuids, {node_uuid})
# Check imported user is part of exported users
self.assertSetEqual(imported_user_emails, set(users_email))
# Check same number of comments (2) pertaining to each user were ex- and imported
self.assertEqual(len(imported_user_one_comment_uuids),
len(user_one_comments_uuid))
self.assertEqual(len(imported_user_two_comment_uuids),
len(user_two_comments_uuid))
# Check imported comments equal exported comments pertaining to specific user
self.assertSetEqual(imported_user_one_comment_uuids,
set(user_one_comments_uuid))
self.assertSetEqual(imported_user_two_comment_uuids,
set(user_two_comments_uuid))
