diff --git a/tests/test_api.py b/tests/test_api.py index 5a2c4e2b..96cc8508 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -2,33 +2,38 @@ import os import sys import pytest -from unittest.mock import patch, Mock -from http.client import HTTPMessage +from unittest.mock import patch, Mock, mock_open thisDir = os.path.dirname(os.path.realpath(__file__)) repoDir = os.path.abspath(os.path.join(thisDir,'../')) sys.path.append(repoDir) -from utils import makeRequestWithRetry +from utils import ( + DEFAULT_REQUEST_TIMEOUT, + UPLOAD_REQUEST_TIMEOUT, + makeRequestWithRetry, + uploadFileToS3, +) logging.getLogger('urllib3').setLevel(logging.DEBUG) @patch("requests.Session.request") -def test_get(mock_response): +def test_get(mock_request): status_code = 200 - mock_response.return_value.status_code = status_code + mock_request.return_value.status_code = status_code response = makeRequestWithRetry('GET', 'https://test.com', retries=2) assert response.status_code == status_code - mock_response.assert_called_once_with('GET', 'https://test.com', - headers=None, - data=None, - params=None, - files=None) + mock_request.assert_called_once_with('GET', 'https://test.com', + headers=None, + data=None, + params=None, + files=None, + timeout=DEFAULT_REQUEST_TIMEOUT) @patch("requests.Session.request") -def test_put(mock_response): +def test_put(mock_request): status_code = 201 - mock_response.return_value.status_code = status_code + mock_request.return_value.status_code = status_code data = { "key1": "value1", @@ -47,20 +52,44 @@ def test_put(mock_response): retries=2) assert response.status_code == status_code - mock_response.assert_called_once_with('POST', + mock_request.assert_called_once_with('POST', 'https://test.com', data=data, headers={"Authorization": "my_token"}, params=params, - files=None) + files=None, + timeout=DEFAULT_REQUEST_TIMEOUT) + +@patch("builtins.open", new_callable=mock_open, read_data=b"file-data") +@patch("utils.makeRequestWithRetry") +def test_upload_timeout(mock_request, mock_file): + mock_request.return_value.json.return_value = { + 'url': 'https://upload.test.com', + 'fields': {'key': 'uploaded-file'}, + } + + key = uploadFileToS3('/tmp/upload.mov') + + assert key == 'uploaded-file' + mock_request.assert_any_call('POST', + 'https://upload.test.com', + data={'key': 'uploaded-file'}, + files={'file': mock_file.return_value}, + timeout=UPLOAD_REQUEST_TIMEOUT) @patch("urllib3.connectionpool.HTTPConnectionPool._get_conn") -def test_success_after_retries(mock_response): - mock_response.return_value.getresponse.side_effect = [ - Mock(status=500, msg=HTTPMessage()), - Mock(status=502, msg=HTTPMessage()), - Mock(status=200, msg=HTTPMessage()), - Mock(status=429, msg=HTTPMessage()), +def test_success_after_retries(mock_get_conn): + def make_response(status): + response = Mock(status=status, headers={}) + response.stream.return_value = [] + response._original_response = None + return response + + mock_get_conn.return_value.getresponse.side_effect = [ + make_response(500), + make_response(502), + make_response(200), + make_response(429), ] response = makeRequestWithRetry('GET', @@ -69,11 +98,11 @@ def test_success_after_retries(mock_response): backoff_factor=0.1) assert response.status_code == 200 - assert mock_response.call_count == 3 + assert mock_get_conn.call_count == 3 # The httpbin test remains commented out for stability reasons # def test_httpbin(): # response = makeRequestWithRetry('GET', # 'https://httpbin.org/status/500', # retries=4, -# backoff_factor=0.1) \ No newline at end of file +# backoff_factor=0.1) diff --git a/utils.py b/utils.py index 962aa74b..7a414c81 100644 --- a/utils.py +++ b/utils.py @@ -28,6 +28,8 @@ API_URL = getAPIURL() API_TOKEN = getToken() +DEFAULT_REQUEST_TIMEOUT = (10, 10) +UPLOAD_REQUEST_TIMEOUT = (10, 300) DEPTH_DB_LEVEL = 10 DEPTH_DB_TRANSFORM = "vertical_delta_shuffle16" DEPTH_CONTAINER_MAGIC = b"OCDEPTHDB1\n" @@ -122,7 +124,8 @@ def uploadFileToS3(filePath): makeRequestWithRetry('POST', r['url'], data=r['fields'], - files=files) + files=files, + timeout=UPLOAD_REQUEST_TIMEOUT) return r['fields']['key'] @@ -2086,7 +2089,8 @@ def postProcessedDuration(trial_url, duration): # utils for common HTTP requests def makeRequestWithRetry(method, url, headers=None, data=None, params=None, files=None, - retries=5, backoff_factor=1): + retries=5, backoff_factor=1, + timeout=DEFAULT_REQUEST_TIMEOUT): """ Makes an HTTP request with retry logic and returns the Response object. @@ -2099,6 +2103,8 @@ def makeRequestWithRetry(method, url, params (dict): URL query parameters. retries (int): Number of retry attempts. backoff_factor (float): Backoff factor for exponential delays. + timeout (float or tuple): Seconds to wait for connection/response + activity, as accepted by requests.Session().request(). Returns: requests.Response: The response object for further processing. @@ -2118,7 +2124,7 @@ def makeRequestWithRetry(method, url, headers=headers, data=data, params=params, - files=files) + files=files, + timeout=timeout) response.raise_for_status() return response -