def test_statement_sensor(mocker):
sen = LivyStatementSensor(
session_id=123,
statement_id=456,
poke_interval=1,
timeout=3,
task_id="test_statement_sensor",
)
response = mock_http_calls(
200,
content=b"""
{
"id": 123,
"state": "available",
"output": {
"status": "ok"
}
}""",
)
mocker.patch.object(HttpHook, "get_conn", return_value=response)
sen.execute({})
def test_statement_sensor_invalid_timings(poke_interval, timeout):
print(f"\n\nRunning test_statement_sensor_invalid_timings with:\n"
f"poke_interval={poke_interval}, timeout={timeout}")
with raises(AirflowException) as ae:
LivyStatementSensor(
session_id=123,
statement_id=456,
poke_interval=poke_interval,
timeout=timeout,
task_id="test_statement_sensor_invalid_timings",
)
print(f"Set up the statement sensor with invalid timings, "
f"got the expected exception:\n<{ae.value}>\n")
def test_statement_sensor_malformed_status(mocker):
sen = LivyStatementSensor(
session_id=123,
statement_id=456,
poke_interval=2,
timeout=5,
task_id="test_statement_sensor_malformed_status",
)
http_response = mock_http_calls(
200,
content=b"""
{
"id": 123,
"state": "available",
"output": {
"nostatus": "here"
}
}""",
)
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowBadRequest) as ae:
sen.execute({})
print(f"Imitated invalid 'output' object for statement sensor, "
f"got the expected exception:\n<{ae.value}>\n")