Add tests and modify dynamic timeout calculation

Author: wojciech-bochniarz-4ss

datareservoirio/storage/storage.py

@@ -363,7 +363,7 @@ def _df_to_blob(df, blob_url, session=_BLOBSTORAGE_SESSION, dynamic_read_timeout
 
 
 def _calculate_timeout(file_size_bytes):
-    bytes_per_second = 1 * 1024 * 1024 # 1MB/s
+    bytes_per_second = 0.5 * 1024 * 1024 # 0,5 MB/s
     min_timeout = 30
-    timeout = max(min_timeout, (file_size_bytes / bytes_per_second) * 1.5)
+    timeout = int(max(min_timeout, (file_size_bytes / bytes_per_second) * 1.5)) # 1.5 tolerance
     return timeout

tests/test_storage/test_storage.py

@@ -5,6 +5,7 @@
 from unittest.mock import ANY, call
 
 import pandas as pd
+import numpy as np
 import pytest
 import requests
 from requests import HTTPError
@@ -16,6 +17,24 @@
 
 TEST_PATH = Path(__file__).parent
 
+def helper_function(func):
+    func._is_helper = True
+    return func
+
+@helper_function
+def generate_dataframe(target_bytes):
+    NUM_COLUMNS = 5
+    BYTES_PER_VALUE = 28  
+    
+    bytes_per_row = NUM_COLUMNS * BYTES_PER_VALUE
+    num_rows = max(1, int(target_bytes / bytes_per_row))
+    
+    data = {
+        f'col_{i}': np.random.randn(num_rows)
+        for i in range(NUM_COLUMNS)
+    }
+    
+    return pd.DataFrame(data)
 
 class Test__blob_to_df:
     """
@@ -556,6 +575,119 @@ def test_put(
 
         mock_requests.assert_has_calls(calls_expected)
 
+    def test_put_dynamic_timeout_sets_correct_value(
+        self,
+        mock_requests,
+        storage_no_cache,
+        response_cases,
+    ):
+        response_cases.set("group4")
+        df = generate_dataframe(20_000_000) # 20 MB
+
+        storage_no_cache.put(
+            df,
+            "http://example/blob/url",
+            (
+                "POST",
+                "https://reservoir-api.4subsea.net/api/files/commit",
+                {"json": {"FileId": "1234"}, "timeout": 10},
+            ),
+            True,
+        )
+
+        calls_expected = [
+            call(
+                method="put",
+                url="http://example/blob/url",
+                headers={"x-ms-blob-type": "BlockBlob"},
+                data=ANY,
+                timeout=(30, 40),
+            ),
+            call(
+                method="POST",
+                url="https://reservoir-api.4subsea.net/api/files/commit",
+                json={"FileId": "1234"},
+                timeout=10,
+            ),
+        ]
+
+        mock_requests.assert_has_calls(calls_expected)
+
+    def test_put_dynamic_timeout_sets_minimal_value(
+        self,
+        mock_requests,
+        storage_no_cache,
+        response_cases,
+    ):
+        response_cases.set("group4")
+        df = generate_dataframe(1_000_000) # 1 MB
+
+        storage_no_cache.put(
+            df,
+            "http://example/blob/url",
+            (
+                "POST",
+                "https://reservoir-api.4subsea.net/api/files/commit",
+                {"json": {"FileId": "1234"}, "timeout": 10},
+            ),
+            True,
+        )
+
+        calls_expected = [
+            call(
+                method="put",
+                url="http://example/blob/url",
+                headers={"x-ms-blob-type": "BlockBlob"},
+                data=ANY,
+                timeout=(30, 30),
+            ),
+            call(
+                method="POST",
+                url="https://reservoir-api.4subsea.net/api/files/commit",
+                json={"FileId": "1234"},
+                timeout=10,
+            ),
+        ]
+
+        mock_requests.assert_has_calls(calls_expected)
+
+    def test_put_without_dynamic_timeout_doesnt_set_its_value(
+        self,
+        mock_requests,
+        storage_no_cache,
+        response_cases,
+    ):
+        response_cases.set("group4")
+        df = generate_dataframe(20_000_000) # 20 MB
+
+        storage_no_cache.put(
+            df,
+            "http://example/blob/url",
+            (
+                "POST",
+                "https://reservoir-api.4subsea.net/api/files/commit",
+                {"json": {"FileId": "1234"}, "timeout": 10},
+            )
+        )
+
+        calls_expected = [
+            call(
+                method="put",
+                url="http://example/blob/url",
+                headers={"x-ms-blob-type": "BlockBlob"},
+                data=ANY,
+                timeout=(30, None),
+            ),
+            call(
+                method="POST",
+                url="https://reservoir-api.4subsea.net/api/files/commit",
+                json={"FileId": "1234"},
+                timeout=10,
+            ),
+        ]
+
+        mock_requests.assert_has_calls(calls_expected)
+
     def test_put_raise_for_status(self, storage_no_cache, data_float):
         df = data_float.as_dataframe()