From e9d4ddfaece229fbd2df1a1bfa3f4f513f33d1ac Mon Sep 17 00:00:00 2001
From: Andy Roxby <107427605+aroxby-wayscript@users.noreply.github.com>
Date: Thu, 16 Feb 2023 10:27:45 -0500
Subject: [PATCH] api: add `one-shot` option to container `stats` (#3089)

Signed-off-by: Andy Roxby <107427605+aroxby-wayscript@users.noreply.github.com>
---
 docker/api/container.py          | 24 ++++++++++++++++++++----
 tests/unit/api_container_test.py | 13 ++++++++++++-
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/docker/api/container.py b/docker/api/container.py
index ce483710c..82a892288 100644
--- a/docker/api/container.py
+++ b/docker/api/container.py
@@ -1126,7 +1126,7 @@ def start(self, container, *args, **kwargs):
         self._raise_for_status(res)
 
     @utils.check_resource('container')
-    def stats(self, container, decode=None, stream=True):
+    def stats(self, container, decode=None, stream=True, one_shot=None):
         """
         Stream statistics for a specific container. Similar to the
         ``docker stats`` command.
@@ -1138,6 +1138,9 @@ def stats(self, container, decode=None, stream=True):
                 False by default.
             stream (bool): If set to false, only the current stats will be
                 returned instead of a stream. True by default.
+            one_shot (bool): If set to true, Only get a single stat instead of
+                waiting for 2 cycles. Must be used with stream=false. False by
+                default.
 
         Raises:
             :py:class:`docker.errors.APIError`
@@ -1145,16 +1148,29 @@ def stats(self, container, decode=None, stream=True):
 
         """
         url = self._url("/containers/{0}/stats", container)
+        params = {
+            'stream': stream
+        }
+        if one_shot is not None:
+            if utils.version_lt(self._version, '1.41'):
+                raise errors.InvalidVersion(
+                    'one_shot is not supported for API version < 1.41'
+                )
+            params['one-shot'] = one_shot
         if stream:
-            return self._stream_helper(self._get(url, stream=True),
+            if one_shot:
+                raise errors.InvalidArgument(
+                    'one_shot is only available in conjunction with '
+                    'stream=False'
+                )
+            return self._stream_helper(self._get(url, params=params),
                                        decode=decode)
         else:
             if decode:
                 raise errors.InvalidArgument(
                     "decode is only available in conjunction with stream=True"
                 )
-            return self._result(self._get(url, params={'stream': False}),
-                                json=True)
+            return self._result(self._get(url, params=params), json=True)
 
     @utils.check_resource('container')
     def stop(self, container, timeout=None):
diff --git a/tests/unit/api_container_test.py b/tests/unit/api_container_test.py
index d7b356c44..c605da371 100644
--- a/tests/unit/api_container_test.py
+++ b/tests/unit/api_container_test.py
@@ -1529,7 +1529,18 @@ def test_container_stats(self):
             'GET',
             url_prefix + 'containers/' + fake_api.FAKE_CONTAINER_ID + '/stats',
             timeout=60,
-            stream=True
+            params={'stream': True}
+        )
+
+    def test_container_stats_with_one_shot(self):
+        self.client.stats(
+            fake_api.FAKE_CONTAINER_ID, stream=False, one_shot=True)
+
+        fake_request.assert_called_with(
+            'GET',
+            url_prefix + 'containers/' + fake_api.FAKE_CONTAINER_ID + '/stats',
+            timeout=60,
+            params={'stream': False, 'one-shot': True}
         )
 
     def test_container_top(self):