[Monitor][Ingestion] Add additional tests (Azure#28746)

pvaneck · web-flow · commit 340d2ff889c4 · 2023-02-14T17:16:00.000-08:00
- Test cases added for aborting an upload via error handler
- Test cases added for &gt; 1MB data uploads

Signed-off-by: Paul Van Eck &lt;paulvaneck@microsoft.com&gt;
diff --git a/sdk/monitor/azure-monitor-ingestion/assets.json b/sdk/monitor/azure-monitor-ingestion/assets.json
@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/monitor/azure-monitor-ingestion",
-  "Tag": "python/monitor/azure-monitor-ingestion_70c5ade383"
+  "Tag": "python/monitor/azure-monitor-ingestion_12276a5674"
 }
diff --git a/sdk/monitor/azure-monitor-ingestion/tests/conftest.py b/sdk/monitor/azure-monitor-ingestion/tests/conftest.py
@@ -23,6 +23,8 @@
 # IN THE SOFTWARE.
 #
 # --------------------------------------------------------------------------
+from datetime import datetime
+
 import pytest
 
 from devtools_testutils.sanitizers import (
@@ -80,3 +82,17 @@ def monitor_info(environment_variables):
         "dce": environment_variables.get(ENV_DCE),
         "dcr_id": environment_variables.get(ENV_DCR_ID)
     }
+
+
+@pytest.fixture(scope="session")
+def large_data():
+    logs = []
+    content = "a" * (1024 * 100) # 100 KiB string
+
+    # Ensure total size is > 2 MiB data
+    for i in range(24):
+        logs.append({
+            "Time": datetime.now().isoformat(),
+            "AdditionalContext": content
+        })
+    return logs
diff --git a/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion.py b/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion.py
@@ -6,6 +6,7 @@
 import gzip
 import json
 import os
+from unittest import mock
 import uuid
 
 import pytest
@@ -43,6 +44,12 @@ def test_send_logs(self, recorded_test, monitor_info):
         with client:
             client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=LOGS_BODY)
 
+    def test_send_logs_large(self, recorded_test, monitor_info, large_data):
+        client = self.create_client_from_credential(
+            LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
+        with client:
+            client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=large_data)
+
     def test_send_logs_error(self, recorded_test, monitor_info):
         client = self.create_client_from_credential(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
@@ -92,3 +99,44 @@ def test_send_logs_gzip_file(self, monitor_info):
         with open(temp_file, 'rb') as f:
             client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=f)
         os.remove(temp_file)
+
+
+    def test_abort_error_handler(self, monitor_info):
+        client = self.create_client_from_credential(
+            LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
+
+        class TestException(Exception):
+            pass
+
+        def on_error(e):
+            on_error.called = True
+            if isinstance(e.error, RuntimeError):
+                raise TestException("Abort")
+            return
+
+        on_error.called = False
+
+        with client:
+            # No exception should be raised
+            with mock.patch("azure.monitor.ingestion._operations._patch.GeneratedOps.upload",
+                            side_effect=ConnectionError):
+                client.upload(
+                    rule_id=monitor_info['dcr_id'],
+                    stream_name=monitor_info['stream_name'],
+                    logs=LOGS_BODY,
+                    on_error=on_error)
+
+            assert on_error.called
+
+            on_error.called = False
+            # Exception should now be raised since error handler checked for RuntimeError.
+            with mock.patch("azure.monitor.ingestion._operations._patch.GeneratedOps.upload",
+                            side_effect=RuntimeError):
+                with pytest.raises(TestException):
+                    client.upload(
+                        rule_id=monitor_info['dcr_id'],
+                        stream_name=monitor_info['stream_name'],
+                        logs=LOGS_BODY,
+                        on_error=on_error)
+
+        assert on_error.called
diff --git a/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion_async.py b/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion_async.py
@@ -6,6 +6,7 @@
 import gzip
 import os
 import json
+from unittest import mock
 import uuid
 
 import pytest
@@ -44,7 +45,17 @@ async def test_send_logs_async(self, recorded_test, monitor_info):
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
         async with client:
             await client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=LOGS_BODY)
-        credential.close()
+        await credential.close()
+
+    @pytest.mark.asyncio
+    async def test_send_logs_large(self, recorded_test, monitor_info, large_data):
+        credential = self.get_credential(LogsIngestionClient, is_async=True)
+        client = self.create_client_from_credential(
+            LogsIngestionClient, credential, endpoint=monitor_info['dce'])
+        async with client:
+            await client.upload(
+                rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=large_data)
+        await credential.close()
 
     @pytest.mark.asyncio
     async def test_send_logs_error(self, recorded_test, monitor_info):
@@ -56,7 +67,7 @@ async def test_send_logs_error(self, recorded_test, monitor_info):
         with pytest.raises(HttpResponseError) as ex:
             async with client:
                 await client.upload(rule_id='bad-rule', stream_name=monitor_info['stream_name'], logs=body)
-        credential.close()
+        await credential.close()
 
     @pytest.mark.asyncio
     async def test_send_logs_error_custom(self, recorded_test, monitor_info):
@@ -76,7 +87,7 @@ async def on_error(e):
             await client.upload(
                 rule_id='bad-rule', stream_name=monitor_info['stream_name'], logs=body, on_error=on_error)
         assert on_error.called
-        credential.close()
+        await credential.close()
 
     @pytest.mark.asyncio
     async def test_send_logs_json_file(self, recorded_test, monitor_info):
@@ -92,7 +103,7 @@ async def test_send_logs_json_file(self, recorded_test, monitor_info):
             with open(temp_file, 'r') as f:
                 await client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=f)
         os.remove(temp_file)
-        credential.close()
+        await credential.close()
 
     @pytest.mark.asyncio
     @pytest.mark.live_test_only("Issues recording binary streams with test-proxy")
@@ -109,4 +120,49 @@ async def test_send_logs_gzip_file(self, monitor_info):
             with open(temp_file, 'rb') as f:
                 await client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=f)
         os.remove(temp_file)
-        credential.close()
+        await credential.close()
+
+
+    @pytest.mark.asyncio
+    async def test_abort_error_handler(self, monitor_info):
+        credential = self.get_credential(LogsIngestionClient, is_async=True)
+        client = self.create_client_from_credential(
+            LogsIngestionClient, credential, endpoint=monitor_info['dce'])
+        body = [{"foo": "bar"}]
+
+        class TestException(Exception):
+            pass
+
+        async def on_error(e):
+            on_error.called = True
+            if isinstance(e.error, RuntimeError):
+                raise TestException("Abort")
+            return
+
+        on_error.called = False
+
+        async with client:
+            # No exception should be raised
+            with mock.patch("azure.monitor.ingestion.aio._operations._patch.GeneratedOps.upload",
+                            side_effect=ConnectionError):
+                await client.upload(
+                    rule_id=monitor_info['dcr_id'],
+                    stream_name=monitor_info['stream_name'],
+                    logs=LOGS_BODY,
+                    on_error=on_error)
+
+            assert on_error.called
+
+            on_error.called = False
+            # Exception should now be raised since error handler checked for RuntimeError.
+            with mock.patch("azure.monitor.ingestion.aio._operations._patch.GeneratedOps.upload",
+                            side_effect=RuntimeError):
+                with pytest.raises(TestException):
+                    await client.upload(
+                        rule_id=monitor_info['dcr_id'],
+                        stream_name=monitor_info['stream_name'],
+                        logs=LOGS_BODY,
+                        on_error=on_error)
+
+            assert on_error.called
+        await credential.close()

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`"AssetsRepo": "Azure/azure-sdk-assets",`
`3`	`3`	`"AssetsRepoPrefixPath": "python",`
`4`	`4`	`"TagPrefix": "python/monitor/azure-monitor-ingestion",`
`5`		`- "Tag": "python/monitor/azure-monitor-ingestion_70c5ade383"`
	`5`	`+ "Tag": "python/monitor/azure-monitor-ingestion_12276a5674"`
`6`	`6`	`}`