test: retry on 5xx codes (#147)

## Description
Test retrying on 5xx codes

Author: mackurzawa

_test_unstructured_client/unit/test_custom_hooks.py

@@ -13,7 +13,8 @@
 FAKE_KEY = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
 
 
-def test_unit_retry_with_backoff_does_retry(caplog):
+@pytest.mark.parametrize("status_code", [500, 503])
+def test_unit_backoff_strategy_logs_retries_5XX(status_code: int, caplog):
     caplog.set_level(logging.INFO)
     filename = "README.md"
     backoff_strategy = BackoffStrategy(
@@ -24,50 +25,54 @@ def test_unit_retry_with_backoff_does_retry(caplog):
     )
 
     with requests_mock.Mocker() as mock:
-        # mock a 502 status code for POST requests to the api
-        mock.post("https://api.unstructuredapp.io/general/v0/general", status_code=502)
+        # mock a 500/503 status code for POST requests to the api
+        mock.post("https://api.unstructuredapp.io/general/v0/general", status_code=status_code)
         session = UnstructuredClient(api_key_auth=FAKE_KEY)
 
         with open(filename, "rb") as f:
             files = shared.Files(content=f.read(), file_name=filename)
 
         req = shared.PartitionParameters(files=files)
+        with pytest.raises(Exception):
+            session.general.partition(req, retries=retries)
 
-        with pytest.raises(Exception) as excinfo:
-            resp = session.general.partition(req, retries=retries)
-            assert resp.status_code == 502
-            assert "API error occurred" in str(excinfo.value)
-
-        # the number of retries varies
-        assert len(mock.request_history) > 1
+    pattern = re.compile(f"Failed to process a request due to API server error with status code {status_code}. "
+                        "Attempting retry number 1 after sleep.")
+    assert bool(pattern.search(caplog.text))
 
 
-@pytest.mark.parametrize("status_code", [500, 503])
-def test_unit_backoff_strategy_logs_retries_5XX(status_code: int, caplog):
-    caplog.set_level(logging.INFO)
+@pytest.mark.parametrize(
+    ("status_code", "expect_retry"),
+    [
+        [500, False],
+        [502, True],
+        [503, True],
+        [504, True],
+    ]
+)
+def test_unit_number_of_retries_in_5xx(status_code: int, expect_retry: bool):
     filename = "README.md"
     backoff_strategy = BackoffStrategy(
-        initial_interval=10, max_interval=100, exponent=1.5, max_elapsed_time=300
+        initial_interval=1, max_interval=10, exponent=1.5, max_elapsed_time=300
     )
     retries = RetryConfig(
         strategy="backoff", backoff=backoff_strategy, retry_connection_errors=True
     )
 
     with requests_mock.Mocker() as mock:
-        # mock a 500/503 status code for POST requests to the api
         mock.post("https://api.unstructuredapp.io/general/v0/general", status_code=status_code)
         session = UnstructuredClient(api_key_auth=FAKE_KEY)
 
         with open(filename, "rb") as f:
             files = shared.Files(content=f.read(), file_name=filename)
 
         req = shared.PartitionParameters(files=files)
-        with pytest.raises(Exception):
+        with pytest.raises(Exception, match=f"unknown content-type received: None: Status {status_code}"):
             session.general.partition(req, retries=retries)
-
-    pattern = re.compile(f"Failed to process a request due to API server error with status code {status_code}. "
-                        "Attempting retry number 1 after sleep.")
-    assert bool(pattern.search(caplog.text))
+    if expect_retry:
+        assert len(mock.request_history) > 1
+    else:
+        assert len(mock.request_history) == 1
 
 
 def test_unit_backoff_strategy_logs_retries_connection_error(caplog):