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