Testing other types of exceptions being retryable

WilliamDee · Dec 23, 2023 · 6c9dffe · 6c9dffe
1 parent 8d61a06
commit 6c9dffe
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/unit/test_retries_connect.py b/tests/unit/test_retries_connect.py
@@ -14,10 +14,10 @@ def creds(self):
         # Create a mock credentials object
         return DuckDBCredentials(
             path="foo.db",
-            retries=Retries(connect_attempts=2)
+            retries=Retries(connect_attempts=2, retryable_exceptions=["IOException", "ArithmeticError"])
         )
 
-    @pytest.mark.parametrize("exception", [None, IOException, ValueError])
+    @pytest.mark.parametrize("exception", [None, IOException, ArithmeticError, ValueError])
     def test_initialize_db(self, creds, exception):
         # Mocking the duckdb.connect method
         with patch('duckdb.connect') as mock_connect:
@@ -30,7 +30,7 @@ def test_initialize_db(self, creds, exception):
             else:
                 # Call the initialize_db method
                 Environment.initialize_db(creds)
-                if exception == IOException:
+                if exception in {IOException, ArithmeticError}:
                     assert mock_connect.call_count == creds.retries.connect_attempts
                 else:
                     mock_connect.assert_called_once_with(creds.path, read_only=False, config={})