[Monitor Query] Fix server timeout live tests (#36916)

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/monitor/azure-monitor-query/tests/test_logs_client.py

@@ -73,15 +73,21 @@ def test_logs_single_query_with_partial_success(self, recorded_test, monitor_inf
     def test_logs_server_timeout(self, recorded_test, monitor_info):
         client = self.get_client(LogsQueryClient, self.get_credential(LogsQueryClient))
 
-        with pytest.raises(HttpResponseError) as e:
+        try:
             response = client.query_workspace(
-                monitor_info['workspace_id'],
+                monitor_info["workspace_id"],
                 "range x from 1 to 1000000000000000 step 1 | count",
                 timespan=None,
                 server_timeout=2,
-                retry_total=0
+                retry_total=0,
             )
-        assert 'Gateway timeout' in e.value.message
+        except HttpResponseError as e:
+            assert "Gateway timeout" in e.message
+        else:
+            # Response an be observed as either 504 response code from the gateway or a partial failure 200 response.
+            assert response.status == LogsQueryStatus.PARTIAL
+            assert "timed out" in str(response.partial_error)
+
 
     @pytest.mark.live_test_only("Issues recording dynamic 'id' values in requests/responses")
     def test_logs_query_batch_default(self, monitor_info):

sdk/monitor/azure-monitor-query/tests/test_logs_client_async.py

@@ -47,20 +47,23 @@ async def test_logs_auth_no_timespan(self, monitor_info):
 
     @pytest.mark.asyncio
     async def test_logs_server_timeout(self, recorded_test, monitor_info):
-        client = self.get_client(
-            LogsQueryClient, self.get_credential(LogsQueryClient, is_async=True))
+        client = self.get_client(LogsQueryClient, self.get_credential(LogsQueryClient, is_async=True))
 
-        with pytest.raises(HttpResponseError) as e:
+        try:
             async with client:
-                await client.query_workspace(
-                    monitor_info['workspace_id'],
+                response = await client.query_workspace(
+                    monitor_info["workspace_id"],
                     "range x from 1 to 1000000000000000 step 1 | count",
                     timespan=None,
                     server_timeout=2,
                     retry_total=0,
                 )
-
-        assert 'Gateway timeout' in e.value.message
+        except HttpResponseError as e:
+            assert "Gateway timeout" in e.message
+        else:
+            # Response an be observed as either 504 response code from the gateway or a partial failure 200 response.
+            assert response.status == LogsQueryStatus.PARTIAL
+            assert "timed out" in str(response.partial_error)
 
     @pytest.mark.asyncio
     async def test_logs_query_batch_raises_on_no_timespan(self, monitor_info):