test(sdk): add tests for Api.upload_file_retry
@@ -3,16 +3,19 @@ import hashlib
import os
import tempfile
from pathlib import Path
-from typing import Callable, Mapping, Optional, Tuple, Type, Union
+from typing import Callable, Mapping, Optional, Sequence, Tuple, Type, Union
from unittest.mock import Mock, call
import pytest
import requests
import responses
+import wandb.errors
from wandb.apis import internal
from wandb.errors import CommError
from wandb.sdk.lib import retry
+from .test_retry import MockTime, mock_time # noqa: F401
+
def test_agent_heartbeat_with_no_agent_id_fails():
a = internal.Api()
@@ -356,3 +359,69 @@ class TestUploadFile:
)
assert check_err(e.value), e.value
+
+
+class TestUploadFileRetry:
+ """Tests the retry logic of upload_file_retry.
+
+ Testing the file-upload logic itself is done in TestUploadFile, above;
+ this class just tests the retry logic (though it does make a couple
+ assumptions about status codes, like "400 isn't retriable, 500 is.")
+ """
+
+ @pytest.mark.parametrize(
+ ["schedule", "num_requests"],
+ [
+ ([200, 0], 1),
+ ([500, 500, 200, 0], 3),
+ ],
+ )
+ def test_stops_after_success(
+ self,
+ some_file: Path,
+ mock_responses: responses.RequestsMock,
+ schedule: Sequence[int],
+ num_requests: int,
+ ):
+ handler = Mock(side_effect=[(status, {}, "") for status in schedule])
+ mock_responses.add_callback("PUT", "http://example.com/upload-dst", handler)
+
+ internal.InternalApi().upload_file_retry(
+ "http://example.com/upload-dst",
+ some_file.open("rb"),
+ )
+
+ assert handler.call_count == num_requests
+
+ def test_stops_after_bad_status(
+ self,
+ some_file: Path,
+ mock_responses: responses.RequestsMock,
+ ):
+ handler = Mock(side_effect=[(400, {}, "")])
+ mock_responses.add_callback("PUT", "http://example.com/upload-dst", handler)
+
+ with pytest.raises(wandb.errors.CommError):
+ internal.InternalApi().upload_file_retry(
+ "http://example.com/upload-dst",
+ some_file.open("rb"),
+ )
+ assert handler.call_count == 1
+
+ def test_stops_after_retry_limit_exceeded(
+ self,
+ some_file: Path,
+ mock_responses: responses.RequestsMock,
+ ):
+ num_retries = 8
+ handler = Mock(return_value=(500, {}, ""))
+ mock_responses.add_callback("PUT", "http://example.com/upload-dst", handler)
+
+ with pytest.raises(wandb.errors.CommError):
+ internal.InternalApi().upload_file_retry(
+ "http://example.com/upload-dst",
+ some_file.open("rb"),
+ num_retries=num_retries,
+ )
+
+ assert handler.call_count == num_retries + 1
