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)