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"
)
def test_submit_batch_malformed_json(dag, mocker):
op = LivyBatchOperator(task_id="test_submit_batch_malformed_json", dag=dag)
http_response = mock_http_calls(201, content=b'{"id":{}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowBadRequest) as bre:
op.submit_batch()
print(f"\n\nImitated malformed JSON response when submitting a batch, "
f"got the expected exception:\n<{bre.value}>")
def test_batch_sensor(mocker):
sen = LivyBatchSensor(batch_id=2, task_id="test_batch_sensor")
http_response = mock_http_calls(
200,
content=b'{"id": 2, "state": "success"}',
)
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
sen.execute({})
def test_submit_batch_string_id(dag, mocker):
op = LivyBatchOperator(task_id="test_submit_batch_string_id", dag=dag)
http_response = mock_http_calls(
201, content=b'{"id":"unexpectedly, a string!"}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowException) as ae:
op.submit_batch()
print(f"\n\nImitated server returning a string for a batch ID, "
f"got the expected exception:\n<{ae.value}>")
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_create_session_get_id(dag, mocker):
op = LivySessionOperator(
statements=[],
task_id="test_create_session_get_id",
dag=dag,
)
http_response = mock_http_calls(201, content=b'{"id": 456}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
op.create_session()
assert op.session_id == 456
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_create_session_malformed_json(dag, mocker):
op = LivySessionOperator(statements=[],
task_id="test_create_session_malformed_json",
dag=dag)
http_response = mock_http_calls(201, content=b'{"id":{}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowBadRequest) as bre:
op.create_session()
print(f"\n\nImitated malformed JSON response when creating a session, "
f"got the expected exception:\n<{bre.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_submit_batch_bad_response_codes(dag, mocker, code):
op = LivyBatchOperator(
task_id=f"test_submit_batch_bad_response_codes_{code}", dag=dag)
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:
op.submit_batch()
print(f"\n\nImitated the {code} error response when submitting a batch, "
f"got the expected exception:\n<{ae.value}>")
def test_create_session_bad_response_codes(dag, mocker, code):
op = LivySessionOperator(statements=[],
task_id="test_create_session_bad_response_codes",
dag=dag)
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:
op.create_session()
print(f"\n\nImitated the {code} error response when creating a session, "
f"got the expected exception:\n<{ae.value}>")
def test_session_creation_sensor_unknown_state(mocker):
sen = LivySessionCreationSensor(
session_id=123,
poke_interval=2,
timeout=5,
task_id="test_session_creation_sensor_malformed_json",
)
http_response = mock_http_calls(200,
content=b'{"id": 123, "state": "IDK"}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowException) as ae:
sen.execute({})
print(f"Imitated unknown session state for session creation sensor, "
f"got the expected exception:\n<{ae.value}>\n")
def test_session_creation_sensor(mocker):
sen = LivySessionCreationSensor(
session_id=123,
poke_interval=1,
timeout=3,
task_id="test_session_creation_sensor",
)
response = mock_http_calls(
200,
content=b'{"id": 123, "state": "idle"}',
)
mocker.patch.object(HttpHook, "get_conn", return_value=response)
sen.execute({})
assert sen.session_id == 123
def test_statement_sensor_malformed_json(mocker):
sen = LivyStatementSensor(
session_id=123,
statement_id=456,
poke_interval=2,
timeout=5,
task_id="test_statement_sensor_malformed_json",
)
http_response = mock_http_calls(200,
content=b'{"id": 123, "state: "running"}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowBadRequest) as ae:
sen.execute({})
print(f"Imitated malformed JSON response for statement sensor, "
f"got the expected exception:\n<{ae.value}>\n")
def test_session_creation_sensor_timeout(mocker):
sen = LivySessionCreationSensor(
session_id=123,
poke_interval=1,
timeout=3,
task_id="test_session_creation_sensor_timeout",
)
http_response = mock_http_calls(
200, content=b'{"id": 123, "state": "starting"}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
with raises(AirflowSensorTimeout) as te:
sen.execute({})
assert 2 <= http_response.send.call_count <= 4
print(f"\n\nImitated period=1sec, timeout=2sec while waiting for session "
f"to be created, got the expected exception:\n<{te.value}>\n"
f"API was polled {http_response.send.call_count} times.")
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_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({})
def test_submit_batch_get_id(dag, mocker):
op = LivyBatchOperator(task_id="test_submit_batch_get_id", dag=dag)
http_response = mock_http_calls(201, content=b'{"id": 123}')
mocker.patch.object(HttpHook, "get_conn", return_value=http_response)
op.submit_batch()
assert op.batch_id == 123
