From f725e7db4108b9c76d11ebb6ae67bd67cdd5bab1 Mon Sep 17 00:00:00 2001
From: Chris Patterson <cpatterson@microsoft.com>
Date: Fri, 11 Mar 2022 12:58:58 -0500
Subject: [PATCH] url_helper: add tuple support for readurl timeout

It may be useful to configure connection timeout and read timeout
separately.  Update readurl() to accept a tuple that is supported
by python requests to configure both.

Signed-off-by: Chris Patterson <cpatterson@microsoft.com>
---
 cloudinit/url_helper.py            |  8 ++-
 tests/unittests/test_url_helper.py | 95 +++++++++++++++++++++++-------
 2 files changed, 80 insertions(+), 23 deletions(-)

diff --git a/cloudinit/url_helper.py b/cloudinit/url_helper.py
index c577e8da3ac..b827eba9140 100644
--- a/cloudinit/url_helper.py
+++ b/cloudinit/url_helper.py
@@ -194,7 +194,8 @@ def readurl(
     :param url: Mandatory url to request.
     :param data: Optional form data to post the URL. Will set request_method
         to 'POST' if present.
-    :param timeout: Timeout in seconds to wait for a response
+    :param timeout: Timeout in seconds to wait for a response. May be a tuple
+        if specifying (connection timeout, read timeout).
     :param retries: Number of times to retry on exception if exception_cb is
         None or exception_cb returns True for the exception caught. Default is
         to fail with 0 retries on exception.
@@ -229,7 +230,10 @@ def readurl(
         request_method = "POST" if data else "GET"
     req_args["method"] = request_method
     if timeout is not None:
-        req_args["timeout"] = max(float(timeout), 0)
+        if isinstance(timeout, tuple):
+            req_args["timeout"] = timeout
+        else:
+            req_args["timeout"] = max(float(timeout), 0)
     if headers_redact is None:
         headers_redact = []
     manual_tries = 1
diff --git a/tests/unittests/test_url_helper.py b/tests/unittests/test_url_helper.py
index 85810e00c7e..059809d9e27 100644
--- a/tests/unittests/test_url_helper.py
+++ b/tests/unittests/test_url_helper.py
@@ -3,6 +3,7 @@
 import logging
 
 import httpretty
+import pytest
 import requests
 
 from cloudinit import util, version
@@ -116,27 +117,6 @@ def test_read_file_or_url_str_from_url_redacts_noheaders(self):
         self.assertNotIn(REDACTED, logs)
         self.assertIn("sekret", logs)
 
-    @mock.patch(M_PATH + "readurl")
-    def test_read_file_or_url_passes_params_to_readurl(self, m_readurl):
-        """read_file_or_url passes all params through to readurl."""
-        url = "http://hostname/path"
-        response = "This is my url content\n"
-        m_readurl.return_value = response
-        params = {
-            "url": url,
-            "timeout": 1,
-            "retries": 2,
-            "headers": {"somehdr": "val"},
-            "data": "data",
-            "sec_between": 1,
-            "ssl_details": {"cert_file": "/path/cert.pem"},
-            "headers_cb": "headers_cb",
-            "exception_cb": "exception_cb",
-        }
-        self.assertEqual(response, read_file_or_url(**params))
-        params.pop("url")  # url is passed in as a positional arg
-        self.assertEqual([mock.call(url, **params)], m_readurl.call_args_list)
-
     def test_wb_read_url_defaults_honored_by_read_file_or_url_callers(self):
         """Readurl param defaults used when unspecified by read_file_or_url
 
@@ -178,6 +158,79 @@ def request(cls, **kwargs):
         self.assertEqual(m_response, response._response)
 
 
+class TestReadFileOrUrlParameters:
+    @mock.patch(M_PATH + "readurl")
+    @pytest.mark.parametrize(
+        "timeout", [1, 1.2, "1", (1, None), (1, 1), (None, None)]
+    )
+    def test_read_file_or_url_passes_params_to_readurl(
+        self, m_readurl, timeout
+    ):
+        """read_file_or_url passes all params through to readurl."""
+        url = "http://hostname/path"
+        response = "This is my url content\n"
+        m_readurl.return_value = response
+        params = {
+            "url": url,
+            "timeout": timeout,
+            "retries": 2,
+            "headers": {"somehdr": "val"},
+            "data": "data",
+            "sec_between": 1,
+            "ssl_details": {"cert_file": "/path/cert.pem"},
+            "headers_cb": "headers_cb",
+            "exception_cb": "exception_cb",
+        }
+
+        assert response == read_file_or_url(**params)
+        params.pop("url")  # url is passed in as a positional arg
+        assert m_readurl.call_args_list == [mock.call(url, **params)]
+
+    @pytest.mark.parametrize(
+        "readurl_timeout,request_timeout",
+        [
+            (-1, 0),
+            ("-1", 0),
+            (None, None),
+            (1, 1.0),
+            (1.2, 1.2),
+            ("1", 1.0),
+            ((1, None), (1, None)),
+            ((1, 1), (1, 1)),
+            ((None, None), (None, None)),
+        ],
+    )
+    def test_readurl_timeout(self, readurl_timeout, request_timeout):
+        url = "http://hostname/path"
+        m_response = mock.MagicMock()
+
+        class FakeSession(requests.Session):
+            @classmethod
+            def request(cls, **kwargs):
+                expected_kwargs = {
+                    "url": url,
+                    "allow_redirects": True,
+                    "method": "GET",
+                    "headers": {
+                        "User-Agent": "Cloud-Init/%s"
+                        % (version.version_string())
+                    },
+                    "timeout": request_timeout,
+                }
+                if request_timeout is None:
+                    expected_kwargs.pop("timeout")
+
+                assert kwargs == expected_kwargs
+                return m_response
+
+        with mock.patch(
+            M_PATH + "requests.Session", side_effect=[FakeSession()]
+        ):
+            response = read_file_or_url(url, timeout=readurl_timeout)
+
+        assert response._response == m_response
+
+
 class TestRetryOnUrlExc(CiTestCase):
     def test_do_not_retry_non_urlerror(self):
         """When exception is not UrlError return False."""