Merge pull request #250 from RazCrimson/fix-container-stream-issues

Fix stream issues in container.top

podman/api/__init__.py

@@ -11,6 +11,7 @@
     prepare_cidr,
     prepare_timestamp,
     stream_frames,
+    stream_helper,
 )
 from podman.api.tar_utils import create_tar, prepare_containerfile, prepare_containerignore
 from .. import version
@@ -58,4 +59,5 @@ def _api_version(release: str, significant: int = 3) -> str:
     'prepare_filters',
     'prepare_timestamp',
     'stream_frames',
+    'stream_helper',
 ]

podman/api/parse_utils.py

@@ -97,3 +97,14 @@ def stream_frames(response: Response) -> Iterator[bytes]:
         if not data:
             return
         yield data
+
+
+def stream_helper(
+    response: Response, decode_to_json: bool = False
+) -> Union[Iterator[bytes], Iterator[Dict[str, Any]]]:
+    """Helper to stream results and optionally decode to json"""
+    for value in response.iter_lines():
+        if decode_to_json:
+            yield json.loads(value)
+        else:
+            yield value

podman/domain/containers.py

@@ -6,7 +6,6 @@
 from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Tuple, Union
 
 import requests
-from requests import Response
 
 from podman import api
 from podman.domain.images import Image
@@ -389,7 +388,9 @@ def start(self, **kwargs) -> None:
         )
         response.raise_for_status()
 
-    def stats(self, **kwargs) -> Iterator[Union[bytes, Dict[str, Any]]]:
+    def stats(
+        self, **kwargs
+    ) -> Union[bytes, Dict[str, Any], Iterator[bytes], Iterator[Dict[str, Any]]]:
         """Return statistics for container.
 
         Keyword Args:
@@ -413,20 +414,9 @@ def stats(self, **kwargs) -> Iterator[Union[bytes, Dict[str, Any]]]:
         response.raise_for_status()
 
         if stream:
-            return self._stats_helper(decode, response.iter_lines())
+            return api.stream_helper(response, decode_to_json=decode)
 
-        return json.loads(response.text) if decode else response.content
-
-    @staticmethod
-    def _stats_helper(
-        decode: bool, body: Iterator[bytes]
-    ) -> Iterator[Union[bytes, Dict[str, Any]]]:
-        """Helper needed to allow stats() to return either a generator or a bytes."""
-        for entry in body:
-            if decode:
-                yield json.loads(entry)
-            else:
-                yield entry
+        return json.loads(response.content) if decode else response.content
 
     def stop(self, **kwargs) -> None:
         """Stop container.
@@ -466,23 +456,20 @@ def top(self, **kwargs) -> Union[Iterator[Dict[str, Any]], Dict[str, Any]]:
             NotFound: when the container no longer exists
             APIError: when the service reports an error
         """
+        stream = kwargs.get("stream", False)
+
         params = {
+            "stream": stream,
             "ps_args": kwargs.get("ps_args"),
-            "stream": kwargs.get("stream", False),
         }
-        response = self.client.get(f"/containers/{self.id}/top", params=params)
+        response = self.client.get(f"/containers/{self.id}/top", params=params, stream=stream)
         response.raise_for_status()
 
-        if params["stream"]:
-            self._top_helper(response)
+        if stream:
+            return api.stream_helper(response, decode_to_json=True)
 
         return response.json()
 
-    @staticmethod
-    def _top_helper(response: Response) -> Iterator[Dict[str, Any]]:
-        for line in response.iter_lines():
-            yield line
-
     def unpause(self) -> None:
         """Unpause processes in container."""
         response = self.client.post(f"/containers/{self.id}/unpause")

podman/tests/unit/test_container.py

@@ -411,6 +411,53 @@ def test_top(self, mock):
         self.assertDictEqual(actual, body)
         self.assertTrue(adapter.called_once)
 
+    @requests_mock.Mocker()
+    def test_top_with_streaming(self, mock):
+        stream = [
+            {
+                "Processes": [
+                    [
+                        'jhonce',
+                        '2417',
+                        '2274',
+                        '0',
+                        'Mar01',
+                        '?',
+                        '00:00:01',
+                        (
+                            '/usr/bin/ssh-agent /bin/sh -c exec -l /bin/bash -c'
+                            ' "/usr/bin/gnome-session"'
+                        ),
+                    ],
+                    ['jhonce', '5544', '3522', '0', 'Mar01', 'pts/1', '00:00:02', '-bash'],
+                    ['jhonce', '6140', '3522', '0', 'Mar01', 'pts/2', '00:00:00', '-bash'],
+                ],
+                "Titles": ["UID", "PID", "PPID", "C", "STIME", "TTY", "TIME CMD"],
+            }
+        ]
+
+        buffer = io.StringIO()
+        for entry in stream:
+            buffer.write(json.JSONEncoder().encode(entry))
+            buffer.write("\n")
+
+        adapter = mock.get(
+            tests.LIBPOD_URL
+            + "/containers/87e1325c82424e49a00abdd4de08009eb76c7de8d228426a9b8af9318ced5ecd/top"
+            "?stream=True",
+            text=buffer.getvalue(),
+        )
+
+        container = Container(attrs=FIRST_CONTAINER, client=self.client.api)
+        top_stats = container.top(stream=True)
+
+        self.assertIsInstance(top_stats, Iterable)
+        for response, actual in zip(top_stats, stream):
+            self.assertIsInstance(response, dict)
+            self.assertDictEqual(response, actual)
+
+        self.assertTrue(adapter.called_once)
+
 
 if __name__ == '__main__':
     unittest.main()

podman/tests/unit/test_parse_utils.py

@@ -1,9 +1,12 @@
 import datetime
 import ipaddress
+import json
 import unittest
-from typing import Any, Optional
-
 from dataclasses import dataclass
+from typing import Any, Iterable, Optional, Tuple
+from unittest import mock
+
+from requests import Response
 
 from podman import api
 
@@ -14,7 +17,7 @@ def test_parse_repository(self):
         class TestCase:
             name: str
             input: Any
-            expected: Optional[str]
+            expected: Tuple[str, Optional[str]]
 
         cases = [
             TestCase(name="empty str", input="", expected=("", None)),
@@ -56,12 +59,36 @@ def test_prepare_timestamp(self):
 
         self.assertEqual(api.prepare_timestamp(None), None)
         with self.assertRaises(ValueError):
-            api.prepare_timestamp("bad input")
+            api.prepare_timestamp("bad input")  # type: ignore
 
     def test_prepare_cidr(self):
         net = ipaddress.IPv4Network("127.0.0.0/24")
         self.assertEqual(api.prepare_cidr(net), ("127.0.0.0", "////AA=="))
 
+    def test_stream_helper(self):
+        streamed_results = [b'{"test":"val1"}', b'{"test":"val2"}']
+        mock_response = mock.Mock(spec=Response)
+        mock_response.iter_lines.return_value = iter(streamed_results)
+
+        streamable = api.stream_helper(mock_response)
+
+        self.assertIsInstance(streamable, Iterable)
+        for expected, actual in zip(streamed_results, streamable):
+            self.assertIsInstance(actual, bytes)
+            self.assertEqual(expected, actual)
+
+    def test_stream_helper_with_decode(self):
+        streamed_results = [b'{"test":"val1"}', b'{"test":"val2"}']
+        mock_response = mock.Mock(spec=Response)
+        mock_response.iter_lines.return_value = iter(streamed_results)
+
+        streamable = api.stream_helper(mock_response, decode_to_json=True)
+
+        self.assertIsInstance(streamable, Iterable)
+        for expected, actual in zip(streamed_results, streamable):
+            self.assertIsInstance(actual, dict)
+            self.assertDictEqual(json.loads(expected), actual)
+
 
 if __name__ == '__main__':
     unittest.main()