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)