def _construct_dag_with_locations(self, source, sink):
# Construct a simple, two-function DAG.
dag = Dag()
dag.name = 'dag'
dag.functions.extend([
Dag.FunctionReference(name=source),
Dag.FunctionReference(name=sink)
])
link = dag.connections.add()
link.source = source
link.sink = sink
# Add the relevant metadata to the policy engine.
source_address = (self.ip, 1)
sink_address = (self.ip, 2)
self.policy.function_locations[source] = {source_address}
self.policy.function_locations[sink] = {sink_address}
return dag, source_address, sink_address
def test_pin_reject(self):
'''
This test explicitly rejects a pin request from the policy and ensures
that it tries again to pin on another node.
'''
# Create two unpinned executors.
address_set = {(self.ip, 1), (self.ip, 2)}
self.policy.unpinned_cpu_executors.update(address_set)
# Create one failing and one successful response.
self.pin_socket.inbox.append(sutils.ok_resp)
self.pin_socket.inbox.append(sutils.error.SerializeToString())
success = self.policy.pin_function(
'dag', Dag.FunctionReference(name='function'), [])
self.assertTrue(success)
# Ensure that both remaining executors have been removed from unpinned
# and that the DAG commit is pending.
self.assertEqual(len(self.policy.unpinned_cpu_executors), 0)
self.assertEqual(len(self.policy.pending_dags), 1)