def test_group_batch_size(self):
"""Test that the group addition in batches works as expected."""
from aiida.orm.groups import Group
# Create 100 nodes
nodes = []
for _ in range(100):
nodes.append(Data().store().backend_entity)
# Add nodes to groups using different batch size. Check in the end the
# correct addition.
batch_sizes = (1, 3, 10, 1000)
for batch_size in batch_sizes:
group = Group(label='test_batches_' + str(batch_size)).store()
group.backend_entity.add_nodes(nodes,
skip_orm=True,
batch_size=batch_size)
self.assertEqual(set(_.pk for _ in nodes),
set(_.pk for _ in group.nodes))
def setUpClass(cls, *args, **kwargs):
super().setUpClass(*args, **kwargs)
from aiida.engine import ProcessState
from aiida.orm.groups import Group
cls.calcs = []
cls.process_label = 'SomeDummyWorkFunctionNode'
# Create 6 WorkFunctionNodes and WorkChainNodes (one for each ProcessState)
for state in ProcessState:
calc = WorkFunctionNode()
calc.set_process_state(state)
# Set the WorkFunctionNode as successful
if state == ProcessState.FINISHED:
calc.set_exit_status(0)
# Give a `process_label` to the `WorkFunctionNodes` so the `--process-label` option can be tested
calc.set_attribute('process_label', cls.process_label)
calc.store()
cls.calcs.append(calc)
calc = WorkChainNode()
calc.set_process_state(state)
# Set the WorkChainNode as failed
if state == ProcessState.FINISHED:
calc.set_exit_status(1)
# Set the waiting work chain as paused as well
if state == ProcessState.WAITING:
calc.pause()
calc.store()
cls.calcs.append(calc)
cls.group = Group('some_group').store()
cls.group.add_nodes(cls.calcs[0])