Streaming support for PodsManager.stats API

podman/domain/pods_manager.py

@@ -1,7 +1,7 @@
 """PodmanResource manager subclassed for Networks."""
 import json
 import logging
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, Iterator
 
 from podman import api
 from podman.domain.manager import Manager
@@ -128,12 +128,14 @@ def remove(self, pod_id: Union[Pod, str], force: Optional[bool] = None) -> None:
         response = self.client.delete(f"/pods/{pod_id}", params={"force": force})
         response.raise_for_status()
 
-    def stats(self, **kwargs) -> Dict[str, Any]:
+    def stats(self, **kwargs) -> Union[List[Dict[str, Any]], Iterator[List[Dict[str, Any]]]]:
         """Resource usage statistics for the containers in pods.
 
         Keyword Args:
             all (bool): Provide statistics for all running pods.
             name (Union[str, List[str]]): Pods to include in report.
+            stream (bool): Stream statistics until cancelled. Default: False.
+            decode (bool): If True, response will be decoded into dict. Default: False.
 
         Raises:
             NotFound: when pod not found
@@ -142,10 +144,20 @@ def stats(self, **kwargs) -> Dict[str, Any]:
         if "all" in kwargs and "name" in kwargs:
             raise ValueError("Keywords 'all' and 'name' are mutually exclusive.")
 
+        # Keeping the default for stream as False to not break existing users
+        # Should probably be changed in a newer major version to match behavior of container.stats
+        stream = kwargs.get("stream", False)
+        decode = kwargs.get("decode", False)
+
         params = {
             "all": kwargs.get("all"),
             "namesOrIDs": kwargs.get("name"),
+            "stream": stream,
         }
-        response = self.client.get("/pods/stats", params=params)
+        response = self.client.get("/pods/stats", params=params, stream=stream)
         response.raise_for_status()
-        return response.json()
+
+        if stream:
+            return api.stream_helper(response, decode_to_json=decode)
+
+        return json.loads(response.content) if decode else response.content

podman/tests/unit/test_podsmanager.py

@@ -1,4 +1,7 @@
+import io
+import json
 import unittest
+from typing import Iterable
 
 import requests_mock
 
@@ -153,10 +156,93 @@ def test_stats(self, mock):
         )
 
         actual = self.client.pods.stats(
-            name="c8b9f5b17dc1406194010c752fc6dcb330192032e27648db9b14060447ecf3b8"
+            name="c8b9f5b17dc1406194010c752fc6dcb330192032e27648db9b14060447ecf3b8",
+        )
+        self.assertEqual(actual, json.dumps(body).encode())
+
+    @requests_mock.Mocker()
+    def test_stats_without_decode(self, mock):
+        body = {
+            "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"],
+        }
+        mock.get(
+            tests.LIBPOD_URL
+            + "/pods/stats"
+            "?namesOrIDs=c8b9f5b17dc1406194010c752fc6dcb330192032e27648db9b14060447ecf3b8",
+            json=body,
+        )
+
+        actual = self.client.pods.stats(
+            name="c8b9f5b17dc1406194010c752fc6dcb330192032e27648db9b14060447ecf3b8", decode=True
         )
         self.assertDictEqual(actual, body)
 
+    @requests_mock.Mocker()
+    def test_top_with_streaming(self, mock):
+        stream = [
+            [
+                {
+                    'CPU': '2.53%',
+                    'MemUsage': '49.15kB / 16.71GB',
+                    'MemUsageBytes': '48KiB / 15.57GiB',
+                    'Mem': '0.00%',
+                    'NetIO': '7.638kB / 430B',
+                    'BlockIO': '-- / --',
+                    'PIDS': '1',
+                    'Pod': '1c948ab42339',
+                    'CID': 'd999c49a7b6c',
+                    'Name': '1c948ab42339-infra',
+                }
+            ],
+            [
+                {
+                    'CPU': '1.46%',
+                    'MemUsage': '57.23B / 16.71GB',
+                    'MemUsageBytes': '48KiB / 15.57GiB',
+                    'Mem': '0.00%',
+                    'NetIO': '7.638kB / 430B',
+                    'BlockIO': '-- / --',
+                    'PIDS': '1',
+                    'Pod': '1c948ab42339',
+                    'CID': 'd999c49a7b6c',
+                    'Name': '1c948ab42339-infra',
+                }
+            ],
+        ]
+
+        buffer = io.StringIO()
+        for entry in stream:
+            buffer.write(json.JSONEncoder().encode(entry))
+            buffer.write("\n")
+
+        adapter = mock.get(
+            tests.LIBPOD_URL + "/pods/stats?stream=True",
+            text=buffer.getvalue(),
+        )
+
+        stream_results = self.client.pods.stats(stream=True, decode=True)
+
+        self.assertIsInstance(stream_results, Iterable)
+        for response, actual in zip(stream_results, stream):
+            self.assertIsInstance(response, list)
+            self.assertListEqual(response, actual)
+
+        self.assertTrue(adapter.called_once)
+
     def test_stats_400(self):
         with self.assertRaises(ValueError):
             self.client.pods.stats(all=True, name="container")