def test_batch_sensor_malformed_json(mocker):
sen = LivyBatchSensor(batch_id=2,
task_id="test_batch_sensor_malformed_json")
http_response = mock_http_calls(200, content='{{"id": 2, "state": }}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowBadRequest) as bre:
sen.poke({})
print(f"\n\nImitated invalid state while waiting for batch to finish, "
f"got the expected exception:\n<{bre.value}>")
def test_batch_sensor_invalid_states(dag, mocker, state):
sen = LivyBatchSensor(batch_id=2,
task_id="test_batch_sensor_invalid_states")
http_response = mock_http_calls(200,
content=f'{{"id": 2, "state": "{state}"}}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowException) as ae:
sen.poke({})
print(f"\n\nImitated invalid state while waiting for batch to finish, "
f"got the expected exception:\n<{ae.value}>")
def test_batch_sensor_bad_response_codes(mocker, code):
sen = LivyBatchSensor(batch_id=2,
task_id="test_batch_sensor_malformed_json")
http_response = mock_http_calls(code,
content=b"Error content",
reason="Good reason")
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowException) as ae:
sen.poke({})
print(f"\n\nImitated the {code} error response when submitting a batch, "
f"got the expected exception:\n<{ae.value}>")
def test_batch_sensor_valid_states(mocker, state):
sen = LivyBatchSensor(batch_id=2, task_id="test_batch_sensor_valid_states")
http_response = mock_http_calls(200,
content=f'{{"id": 2, "state": "{state}"}}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
assert not sen.poke({})