Merge pull request #83 from gthomas-slack/add-statement-progress

Add statement progress to Statement model
acroz · Jun 12, 2020 · 8169136 · 8169136
2 parents 6c0a960 + 44cbcaf
commit 8169136
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 4 deletions.
diff --git a/livy/models.py b/livy/models.py
@@ -112,15 +112,21 @@ class Statement:
     statement_id: int
     state: StatementState
     output: Optional[Output]
+    progress: Optional[float]
 
     @classmethod
     def from_json(cls, session_id: int, data: dict) -> "Statement":
         if data["output"] is None:
             output = None
         else:
             output = Output.from_json(data["output"])
+
         return cls(
-            session_id, data["id"], StatementState(data["state"]), output
+            session_id,
+            data["id"],
+            StatementState(data["state"]),
+            output,
+            data.get("progress"),
         )
 
 

diff --git a/tests/test_models.py b/tests/test_models.py
@@ -58,10 +58,19 @@ def test_session_from_json():
 def test_statement_from_json_no_output():
 
     session_id = 5
-    statement_json = {"id": 10, "state": "running", "output": None}
+    statement_json = {
+        "id": 10,
+        "state": "running",
+        "output": None,
+        "progress": 0.0,
+    }
 
     expected = Statement(
-        session_id, statement_id=10, state=StatementState.RUNNING, output=None
+        session_id,
+        statement_id=10,
+        state=StatementState.RUNNING,
+        output=None,
+        progress=0.0,
     )
 
     assert Statement.from_json(session_id, statement_json) == expected
@@ -72,19 +81,71 @@ 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"}
+    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
     Output.from_json.assert_called_once_with("dummy output")
 
 
+def test_statement_from_json_no_progress(mocker):
+
+    mocker.patch.object(Output, "from_json")
+
+    session_id = 5
+    statement_json = {
+        "id": 10,
+        "state": "running",
+        "output": "dummy output",
+        "progress": None,
+    }
+
+    expected = Statement(
+        session_id,
+        statement_id=10,
+        state=StatementState.RUNNING,
+        output=Output.from_json.return_value,
+        progress=None,
+    )
+
+    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 test_output_textdata_from_json():
 
     output_json = {"status": "ok", "data": {"text/plain": "some output"}}