Skip to content

Commit 47338a5

Browse files
authored
Refactored storage.py put methods (#42502)
* Refactored storage.py and modified put methods of LocalFileBlob and LocalFileStorage * Updated CHANGELOG * Addressed feedback
1 parent 9869606 commit 47338a5

File tree

5 files changed

+104
-26
lines changed

5 files changed

+104
-26
lines changed

sdk/monitor/azure-monitor-opentelemetry-exporter/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
### Features Added
66
- Customer Facing Statsbeat: Added remaining drop codes to base
77
([#42382](https://github.com/Azure/azure-sdk-for-python/pull/42382))
8+
- Refactored the put methods in storage.py for LocalFileBlob and LocalFileStorage
9+
([#42502](https://github.com/Azure/azure-sdk-for-python/pull/42502))
810

911
### Breaking Changes
1012

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/_storage.py

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import random
99
import subprocess
1010
import errno
11+
from typing import Union
1112
from enum import Enum
1213

1314
from azure.monitor.opentelemetry.exporter._utils import PeriodicTask
@@ -37,6 +38,7 @@ def _seconds(seconds):
3738
return datetime.timedelta(seconds=seconds)
3839

3940
class StorageExportResult(Enum):
41+
LOCAL_FILE_BLOB_SUCCESS = 0
4042
CLIENT_STORAGE_DISABLED = 1
4143
CLIENT_PERSISTENCE_CAPACITY_REACHED = 2
4244
CLIENT_READONLY = 3
@@ -60,9 +62,7 @@ def get(self):
6062
pass # keep silent
6163
return None
6264

63-
def put(self, data, lease_period=0):
64-
#TODO: Modify method to remove the return of self as it is not being used anywhere.
65-
# Add typing to method
65+
def put(self, data, lease_period=0) -> Union[StorageExportResult, str]:
6666
try:
6767
fullpath = self.fullpath + ".tmp"
6868
with open(fullpath, "w", encoding="utf-8") as file:
@@ -76,7 +76,7 @@ def put(self, data, lease_period=0):
7676
timestamp = _now() + _seconds(lease_period)
7777
self.fullpath += "@{}.lock".format(_fmt(timestamp))
7878
os.rename(fullpath, self.fullpath)
79-
return self
79+
return StorageExportResult.LOCAL_FILE_BLOB_SUCCESS
8080
except Exception as ex:
8181
return str(ex)
8282

@@ -194,10 +194,7 @@ def get(self):
194194
pass
195195
return None
196196

197-
def put(self, data, lease_period=None):
198-
# TODO: Remove the blob.put result as we are not using it anywhere and use StorageExportResult instead,
199-
# Should still capture exceptions returned from LocalFileBlob.put
200-
# Add typing for method
197+
def put(self, data, lease_period=None) -> Union[StorageExportResult, str]:
201198
try:
202199
if not self._enabled:
203200
if get_local_storage_setup_state_readonly():

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/export/_base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ def _handle_transmit_from_storage(self, envelopes: List[TelemetryItem], result:
223223
# For any exceptions occurred in put method of either LocalFileStorage or LocalFileBlob, track dropped item with reason
224224
_track_dropped_items(self._customer_statsbeat_metrics, envelopes, DropCode.CLIENT_EXCEPTION, result_from_storage_put)
225225
else:
226-
# LocalFileBlob.put returns either an exception(failure, handled above) or the file path(success), eventually that will be removed since this value is not being utilized elsewhere
226+
# LocalFileBlob.put returns StorageExportResult.LOCAL_FILE_BLOB_SUCCESS here. Don't need to track anything in this case.
227227
pass
228228
elif result == ExportResult.SUCCESS:
229229
# Try to send any cached events

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/test_base_exporter.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1066,6 +1066,32 @@ def invalid_get_token():
10661066
ValueError, _get_auth_policy, credential=InvalidTestCredential(), default_auth_policy=TEST_AUTH_POLICY
10671067
)
10681068

1069+
@mock.patch("azure.monitor.opentelemetry.exporter.statsbeat._utils._track_dropped_items")
1070+
@mock.patch("azure.monitor.opentelemetry.exporter.statsbeat._utils._track_dropped_items")
1071+
def test_handle_transmit_from_storage_success_result(self, mock_track_dropped1, mock_track_dropped2):
1072+
"""Test that when storage.put() returns StorageExportResult.LOCAL_FILE_BLOB_SUCCESS,
1073+
the method continues without any special handling."""
1074+
exporter = BaseExporter(disable_offline_storage=False)
1075+
mock_customer_statsbeat = mock.Mock()
1076+
exporter._customer_statsbeat_metrics = mock_customer_statsbeat
1077+
exporter._should_collect_customer_statsbeat = mock.Mock(return_value=True)
1078+
1079+
# Mock storage.put() to return success
1080+
exporter.storage = mock.Mock()
1081+
exporter.storage.put.return_value = StorageExportResult.LOCAL_FILE_BLOB_SUCCESS
1082+
1083+
test_envelopes = [TelemetryItem(name="test", time=datetime.now())]
1084+
serialized_envelopes = [envelope.as_dict() for envelope in test_envelopes]
1085+
exporter._handle_transmit_from_storage(test_envelopes, ExportResult.FAILED_RETRYABLE)
1086+
1087+
# Verify storage.put was called with the serialized envelopes
1088+
exporter.storage.put.assert_called_once_with(serialized_envelopes)
1089+
# Verify that no dropped items were tracked (since it was a success)
1090+
mock_track_dropped1.assert_not_called()
1091+
mock_track_dropped2.assert_not_called()
1092+
# Verify that the customer statsbeat wasn't invoked
1093+
mock_customer_statsbeat.assert_not_called()
1094+
10691095
def test_get_auth_policy_audience(self):
10701096
class TestCredential:
10711097
def get_token():
@@ -2162,6 +2188,36 @@ def test_handle_transmit_from_storage_success_path_validation(self, mock_track_d
21622188
},
21632189
)
21642190
@mock.patch('azure.monitor.opentelemetry.exporter.export._base._track_dropped_items')
2191+
@mock.patch.dict(
2192+
os.environ,
2193+
{
2194+
"APPLICATIONINSIGHTS_STATSBEAT_DISABLED_ALL": "true",
2195+
"APPLICATIONINSIGHTS_STATSBEAT_ENABLED_PREVIEW": "true",
2196+
},
2197+
)
2198+
@mock.patch('azure.monitor.opentelemetry.exporter.export._base._track_dropped_items')
2199+
def test_handle_transmit_from_storage_unexpected_return_value(self, mock_track_dropped1, mock_track_dropped2):
2200+
"""Test that when storage.put() returns an unexpected value type (not StorageExportResult or str),
2201+
the method continues without any special handling."""
2202+
exporter = BaseExporter(disable_offline_storage=False)
2203+
mock_customer_statsbeat = mock.Mock()
2204+
exporter._customer_statsbeat_metrics = mock_customer_statsbeat
2205+
exporter._should_collect_customer_statsbeat = mock.Mock(return_value=True)
2206+
2207+
# Mock storage.put() to return an unexpected value type (int)
2208+
exporter.storage = mock.Mock()
2209+
exporter.storage.put.return_value = 42 # Neither StorageExportResult nor str
2210+
2211+
test_envelopes = [TelemetryItem(name="test", time=datetime.now())]
2212+
exporter._handle_transmit_from_storage(test_envelopes, ExportResult.FAILED_RETRYABLE)
2213+
2214+
# Verify that no dropped items were tracked (since return value isn't handled)
2215+
mock_track_dropped1.assert_not_called()
2216+
mock_track_dropped2.assert_not_called()
2217+
# Verify that the customer statsbeat wasn't invoked
2218+
mock_customer_statsbeat.assert_not_called()
2219+
2220+
@mock.patch("azure.monitor.opentelemetry.exporter.export._base._track_dropped_items")
21652221
def test_handle_transmit_from_storage_string_return_values_trigger_exception_tracking(self, mock_track_dropped):
21662222
"""Test that string return values from storage.put() trigger CLIENT_EXCEPTION tracking"""
21672223
# Save original state

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/test_storage.py

Lines changed: 40 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -70,9 +70,9 @@ def test_put_success_returns_self(self):
7070
test_input = [1, 2, 3]
7171
result = blob.put(test_input)
7272
# Should return the blob itself (self) on success
73-
self.assertIsInstance(result, LocalFileBlob)
74-
self.assertEqual(result, blob)
75-
73+
self.assertIsInstance(result, StorageExportResult)
74+
self.assertEqual(result, StorageExportResult.LOCAL_FILE_BLOB_SUCCESS)
75+
7676
def test_put_file_write_error_returns_string(self):
7777
blob = LocalFileBlob(os.path.join(TEST_FOLDER, "write_error_blob"))
7878
test_input = [1, 2, 3]
@@ -129,8 +129,8 @@ def test_put_with_lease_period_success(self):
129129
lease_period = 60
130130

131131
result = blob.put(test_input, lease_period=lease_period)
132-
self.assertIsInstance(result, LocalFileBlob)
133-
self.assertEqual(result, blob)
132+
self.assertIsInstance(result, StorageExportResult)
133+
self.assertEqual(result, StorageExportResult.LOCAL_FILE_BLOB_SUCCESS)
134134
# File should have .lock extension due to lease period
135135
self.assertTrue(blob.fullpath.endswith(".lock"))
136136

@@ -150,17 +150,17 @@ def test_put_empty_data_success(self):
150150
empty_data = []
151151

152152
result = blob.put(empty_data)
153-
self.assertIsInstance(result, LocalFileBlob)
154-
self.assertEqual(result, blob)
153+
self.assertIsInstance(result, StorageExportResult)
154+
self.assertEqual(result, StorageExportResult.LOCAL_FILE_BLOB_SUCCESS)
155155

156156
def test_put_large_data_success(self):
157157
blob = LocalFileBlob(os.path.join(TEST_FOLDER, "large_data_blob"))
158158
# Create a large list of data
159159
large_data = [{"id": i, "value": f"data_{i}"} for i in range(1000)]
160160

161161
result = blob.put(large_data)
162-
self.assertIsInstance(result, LocalFileBlob)
163-
self.assertEqual(result, blob)
162+
self.assertIsInstance(result, StorageExportResult)
163+
self.assertEqual(result, StorageExportResult.LOCAL_FILE_BLOB_SUCCESS)
164164

165165
# Verify data can be retrieved
166166
retrieved_data = blob.get()
@@ -174,13 +174,25 @@ def test_put_return_type_consistency(self):
174174

175175
# Test successful case
176176
result_success = blob.put(test_input)
177-
self.assertTrue(isinstance(result_success, LocalFileBlob) or isinstance(result_success, str))
177+
self.assertTrue(isinstance(result_success, StorageExportResult) or isinstance(result_success, str))
178178

179179
# Test error case
180180
blob2 = LocalFileBlob(os.path.join(TEST_FOLDER, "consistency_blob2"))
181181
with mock.patch("os.rename", side_effect=Exception("Test error")):
182182
result_error = blob2.put(test_input)
183183
self.assertIsInstance(result_error, str)
184+
185+
def test_put_invalid_return_type(self):
186+
blob = LocalFileBlob(os.path.join(TEST_FOLDER, "invalid_return_blob"))
187+
test_input = [1, 2, 3]
188+
189+
# This tests that even if os.rename somehow returns something unexpected,
190+
# the put method still maintains its type contract
191+
with mock.patch("os.rename", return_value=42):
192+
result = blob.put(test_input)
193+
# Should either convert to string or return StorageExportResult
194+
self.assertTrue(isinstance(result, (StorageExportResult, str)),
195+
f"Expected StorageExportResult or str, got {type(result)}")
184196

185197
@unittest.skip("transient storage")
186198
def test_put(self):
@@ -334,7 +346,7 @@ def test_put_success_returns_localfileblob(self):
334346
test_input = (1, 2, 3)
335347
with LocalFileStorage(os.path.join(TEST_FOLDER, "success_test")) as stor:
336348
result = stor.put(test_input, lease_period=0) # No lease period so file is immediately available
337-
self.assertIsInstance(result, LocalFileBlob)
349+
self.assertIsInstance(result, StorageExportResult)
338350
self.assertEqual(stor.get().get(), test_input)
339351

340352
def test_put_blob_put_failure_returns_string(self):
@@ -379,19 +391,19 @@ def test_put_with_lease_period(self):
379391

380392
with LocalFileStorage(os.path.join(TEST_FOLDER, "lease_test")) as stor:
381393
result = stor.put(test_input, lease_period=custom_lease_period)
382-
self.assertIsInstance(result, LocalFileBlob)
394+
self.assertIsInstance(result, StorageExportResult)
383395
# Verify the file was created with lease period
384-
self.assertTrue(result.fullpath.endswith(".lock"))
396+
self.assertEqual(result, StorageExportResult.LOCAL_FILE_BLOB_SUCCESS)
385397

386398
def test_put_default_lease_period(self):
387399
test_input = (1, 2, 3)
388400

389401
with LocalFileStorage(os.path.join(TEST_FOLDER, "default_lease_test"), lease_period=90) as stor:
390402
result = stor.put(test_input)
391-
self.assertIsInstance(result, LocalFileBlob)
403+
self.assertIsInstance(result, StorageExportResult)
392404
# File should be created with lease (since default lease_period > 0)
393-
self.assertTrue(result.fullpath.endswith(".lock"))
394-
405+
self.assertEqual(result, StorageExportResult.LOCAL_FILE_BLOB_SUCCESS)
406+
395407
def test_check_and_set_folder_permissions_oserror_sets_exception_state(self):
396408
test_input = (1, 2, 3)
397409
test_error_message = "OSError: Permission denied creating directory"
@@ -671,7 +683,7 @@ def test_exception_state_cleared_and_storage_recovery(self):
671683
with LocalFileStorage(os.path.join(TEST_FOLDER, "recovery_test2")) as stor2:
672684
if stor2._enabled: # Storage should be enabled now
673685
result = stor2.put(test_input, lease_period=0)
674-
self.assertIsInstance(result, LocalFileBlob)
686+
self.assertIsInstance(result, StorageExportResult)
675687
retrieved_data = stor2.get().get()
676688
self.assertEqual(retrieved_data, test_input)
677689

@@ -740,6 +752,17 @@ def test_readonly_state_interaction_with_storage_put_method(self):
740752
# When storage is disabled and readonly state is set, put() should return CLIENT_READONLY
741753
result = stor.put(test_input)
742754
self.assertEqual(result, StorageExportResult.CLIENT_READONLY)
755+
756+
def test_storage_put_invalid_return_type(self):
757+
test_input = (1, 2, 3)
758+
759+
with LocalFileStorage(os.path.join(TEST_FOLDER, "invalid_return_test")) as stor:
760+
# Mock _check_storage_size to return a non-boolean value
761+
with mock.patch.object(stor, '_check_storage_size', return_value=42):
762+
result = stor.put(test_input)
763+
# Should maintain return type contract despite invalid internal return
764+
self.assertTrue(isinstance(result, (StorageExportResult, str)),
765+
f"Expected StorageExportResult or str, got {type(result)}")
743766

744767
def test_readonly_state_priority_over_exception_state(self):
745768
test_input = (1, 2, 3)

0 commit comments

Comments
 (0)