def test_statement_from_json_no_output():
session_id = 5
statement_json = {"id": 10, "state": "running", "output": None}
expected = Statement(
session_id, statement_id=10, state=StatementState.RUNNING, output=None
)
assert Statement.from_json(session_id, statement_json) == expected
def test_statement_from_json_no_output():
session_id = 5
statement_json = {'id': 10, 'state': 'running', 'output': None}
expected = Statement(session_id,
statement_id=10,
state=StatementState.RUNNING,
output=None)
assert Statement.from_json(session_id, statement_json) == expected
def test_statement_from_json_with_output(mocker):
mocker.patch.object(Output, 'from_json')
session_id = 5
statement_json = {'id': 10, 'state': 'running', 'output': 'dummy output'}
expected = Statement(session_id,
statement_id=10,
state=StatementState.RUNNING,
output=Output.from_json.return_value)
assert Statement.from_json(session_id, statement_json) == expected
Output.from_json.assert_called_once_with('dummy output')
def test_statement_from_json_with_output(mocker):
mocker.patch.object(Output, "from_json")
session_id = 5
statement_json = {"id": 10, "state": "running", "output": "dummy output"}
expected = Statement(
session_id,
statement_id=10,
state=StatementState.RUNNING,
output=Output.from_json.return_value,
)
assert Statement.from_json(session_id, statement_json) == expected
Output.from_json.assert_called_once_with("dummy output")
def get_statement(self, session_id: int, statement_id: int) -> Statement:
"""Get information about a statement in a session.
:param session_id: The ID of the session.
:param statement_id: The ID of the statement.
"""
response = self._client.get(
f"/sessions/{session_id}/statements/{statement_id}")
return Statement.from_json(session_id, response)
def list_statements(self, session_id: int) -> List[Statement]:
"""Get all the statements in a session.
:param session_id: The ID of the session.
"""
response = self._client.get(f"/sessions/{session_id}/statements")
return [
Statement.from_json(session_id, data)
for data in response["statements"]
]
def test_statement_from_json_no_output():
session_id = 5
statement_json = {
"id": 10,
"state": "running",
"code": "dummy code",
"output": None,
"progress": 0.0,
}
expected = Statement(
session_id,
statement_id=10,
state=StatementState.RUNNING,
code="dummy code",
output=None,
progress=0.0,
)
assert Statement.from_json(session_id, statement_json) == expected
def test_statement_from_json_with_progress(mocker):
mocker.patch.object(Output, "from_json")
session_id = 5
statement_json = {
"id": 10,
"state": "running",
"output": "dummy output",
"progress": 0.5,
}
expected = Statement(
session_id,
statement_id=10,
state=StatementState.RUNNING,
output=Output.from_json.return_value,
progress=0.5,
)
assert Statement.from_json(session_id, statement_json) == expected
def create_statement(self,
session_id: int,
code: str,
kind: StatementKind = None) -> Statement:
data = {"code": code}
if kind is not None:
if self.legacy_server():
LOGGER.warning("statement kind ignored on Livy<0.5.0")
data["kind"] = kind.value
response = self._client.post(f"/sessions/{session_id}/statements",
data=data)
return Statement.from_json(session_id, response)
def create_statement(self,
session_id: int,
code: str,
kind: StatementKind = None) -> Statement:
"""Run a statement in a session.
:param session_id: The ID of the session.
:param code: The code to execute.
:param kind: The kind of code to execute.
"""
data = {"code": code}
if kind is not None:
if self.legacy_server():
LOGGER.warning("statement kind ignored on Livy<0.5.0")
data["kind"] = kind.value
response = self._client.post(f"/sessions/{session_id}/statements",
data=data)
return Statement.from_json(session_id, response)
def get_statement(self, session_id: int, statement_id: int) -> Statement:
response = self._client.get(
f"/sessions/{session_id}/statements/{statement_id}")
return Statement.from_json(session_id, response)
def list_statements(self, session_id: int) -> List[Statement]:
response = self._client.get(f"/sessions/{session_id}/statements")
return [
Statement.from_json(session_id, data)
for data in response["statements"]
]