Fixup unit tests

Author: Adrian Chang

libs/labelbox/tests/unit/export_task/test_export_task.py

@@ -1,4 +1,5 @@
 import json
+import pytest
 
 from unittest.mock import MagicMock, patch
 from labelbox.schema.export_task import ExportTask
@@ -60,6 +61,76 @@ def test_export_task(self):
             export_task = ExportTask(mock_task, is_export_v2=True)
             assert export_task.result[0] == data
 
+    def test_get_buffered_stream_complete(self):
+        with pytest.raises(ExportTask.ExportTaskException):
+            mock_task = MagicMock()
+            mock_task.status = "FAILED"
+            export_task = ExportTask(mock_task, is_export_v2=True)
+            export_task.get_buffered_stream()
+
+    def test_get_buffered_stream_failed(self):
+        with pytest.raises(ExportTask.ExportTaskException):
+            mock_task = MagicMock()
+            mock_task.status = "INPROGRESS"
+            export_task = ExportTask(mock_task, is_export_v2=True)
+            export_task.get_buffered_stream()
+
+    def test_get_buffered_stream(self):
+        with patch('requests.get') as mock_requests_get:
+            mock_task = MagicMock()
+            mock_task.client.execute.side_effect = [
+                {
+                    "task": {
+                        "exportMetadataHeader": {
+                            "total_size": 1,
+                            "total_lines": 1,
+                            "lines": {
+                                "start": 0,
+                                "end": 1
+                            },
+                            "offsets": {
+                                "start": 0,
+                                "end": 0
+                            },
+                            "file": "file"
+                        }
+                    }
+                },
+                {
+                    "task": {
+                        "exportFileFromOffset": {
+                            "total_size": 1,
+                            "total_lines": 1,
+                            "lines": {
+                                "start": 0,
+                                "end": 1
+                            },
+                            "offsets": {
+                                "start": 0,
+                                "end": 0
+                            },
+                            "file": "file"
+                        }
+                    }
+                },
+            ]
+            mock_task.status = "COMPLETE"
+            data = {
+                "data_row": {
+                    "raw_data":
+                        """
+                    {"raw_text":"}{"}
+                    {"raw_text":"\\nbad"}
+                    """
+                }
+            }
+            mock_requests_get.return_value.text = json.dumps(data)
+            mock_requests_get.return_value.content = "b"
+            export_task = ExportTask(mock_task, is_export_v2=True)
+            output_data = []
+            export_task.get_buffered_stream().start(stream_handler=lambda x: output_data.append(x.json))
+            assert data == output_data[0]
+
     def test_export_task_bad_offsets(self):
         with patch('requests.get') as mock_requests_get:
             mock_task = MagicMock()