def test_command_required(self): """ Test that command is required for execution """ op = FabricOperator(task_id=TEST_TASK_ID, fabric_hook=self.hook) with self.assertRaises(AirflowException): op.execute(context={})
def test_ssh_conn_id_or_fabric_hook_required(self): """ Test that either ssh_conn_id or fabric_hook is required for execution """ op = FabricOperator(task_id=TEST_TASK_ID, command="ls") with self.assertRaises(AirflowException) as assertion: op.execute(context={}) self.assertIn("ssh_conn_id", str(assertion.exception)) self.assertIn("fabric_hook", str(assertion.exception))
def test_xcom_push(self): """ Test that an XCom is pushed with the specified name as key and the stdout as value if `xcom_push_key` is set """ task_inst = Mock() op = FabricOperator(task_id=TEST_TASK_ID, fabric_hook=self.hook, command="ls", xcom_push_key="test_xcom") op.execute(context={"task_instance": task_inst}) task_inst.xcom_push.assert_called_with("test_xcom", self.hook.stdout)
def test_fabric_operator_execute_non_zero_exit(self): """ Test that execution with a non-zero exit code fails and the exit code is logged """ self.hook.exit_code = faker.pyint(min_value=1) op = FabricOperator(task_id=TEST_TASK_ID, fabric_hook=self.hook, command="ls") with self.assertRaises(AirflowException) as assertion: op.execute(context={}) self.assertIn(self.hook.exit_code, str(assertion.exception))
def test_fabric_operator_execute_success(self): """ Test that execution with a zero exit code succeeds and called connection.run with the correct parameters """ command = faker.text() env = faker.pydict() pty = faker.pybool() op = FabricOperator(task_id=TEST_TASK_ID, fabric_hook=self.hook, command=command, environment=env, get_pty=pty) self.assertTrue(op.execute(context={})) res = op.execute_fabric_command() res.conn.run.assert_called_with(command=command, pty=pty, env=env, watchers=[], warn=True)