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", "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 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 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 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 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 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"] ]