Fix Azure.Monitor.Query live tests (Azure#21609)

Author: pakrym

sdk/monitor/Azure.Monitor.Query/README.md

@@ -269,7 +269,6 @@ try
 catch (Exception e)
 {
     Console.WriteLine(e);
-    throw;
 }
 ```
 

sdk/monitor/Azure.Monitor.Query/src/Generated/Models/LogsBatchQueryResultInternal.Serialization.cs

@@ -17,6 +17,7 @@ namespace Azure.Monitor.Query
     /// </summary>
     public class LogsQueryClient
     {
+        private static readonly TimeSpan _networkTimeoutOffset = TimeSpan.FromSeconds(15);
         private readonly QueryRestClient _queryClient;
         private readonly ClientDiagnostics _clientDiagnostics;
         private readonly HttpPipeline _pipeline;
@@ -226,7 +227,8 @@ private async Task<Response<LogsQueryResult>> ExecuteAsync(string workspaceId, s
 
             if (options?.ServerTimeout != null)
             {
-                message.NetworkTimeout = options.ServerTimeout;
+                // Offset the service timeout a bit to make sure we have time to receive the response.
+                message.NetworkTimeout = options.ServerTimeout.Value.Add(_networkTimeoutOffset);
             }
 
             if (async)

sdk/monitor/Azure.Monitor.Query/src/LogsQueryClient.cs

@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Text.Json;
+using Azure.Core;
+
+namespace Azure.Monitor.Query.Models
+{
+    [CodeGenModel("logQueryResult")]
+    internal partial class LogsBatchQueryResultInternal: LogsQueryResult
+    {
+        internal ErrorInfo Error { get; }
+
+        // TODO, remove after https://github.com/Azure/azure-sdk-for-net/issues/21655 is fixed
+        internal static LogsBatchQueryResultInternal DeserializeLogsBatchQueryResultInternal(JsonElement element)
+        {
+            Optional<ErrorInfo> error = default;
+            IReadOnlyList<LogsQueryResultTable> tables = default;
+            Optional<JsonElement> statistics = default;
+
+            // This is the workaround to remove the double-encoding
+            if (element.ValueKind == JsonValueKind.String)
+            {
+                try
+                {
+                    using var document = JsonDocument.Parse(element.GetString());
+                    element = document.RootElement.Clone();
+                }
+                catch
+                {
+                    // ignore
+                }
+            }
+
+            foreach (var property in element.EnumerateObject())
+            {
+                if (property.NameEquals("error"))
+                {
+                    if (property.Value.ValueKind == JsonValueKind.Null)
+                    {
+                        property.ThrowNonNullablePropertyIsNull();
+                        continue;
+                    }
+                    error = ErrorInfo.DeserializeErrorInfo(property.Value);
+                    continue;
+                }
+                if (property.NameEquals("tables"))
+                {
+                    List<LogsQueryResultTable> array = new List<LogsQueryResultTable>();
+                    foreach (var item in property.Value.EnumerateArray())
+                    {
+                        array.Add(LogsQueryResultTable.DeserializeLogsQueryResultTable(item));
+                    }
+                    tables = array;
+                    continue;
+                }
+                if (property.NameEquals("statistics"))
+                {
+                    statistics = property.Value.Clone();
+                    continue;
+                }
+            }
+            return new LogsBatchQueryResultInternal(tables, statistics, error.Value);
+        }
+    }
+}

sdk/monitor/Azure.Monitor.Query/src/Models/LogsBatchQueryResultInternal.cs

@@ -217,7 +217,6 @@ public async Task BadRequest()
             catch (Exception e)
             {
                 Console.WriteLine(e);
-                throw;
             }
             #endregion
         }

sdk/monitor/Azure.Monitor.Query/src/Models/LogsQueryResult.cs

@@ -32,7 +32,10 @@ private LogsQueryClient CreateClient()
             return InstrumentClient(new LogsQueryClient(
                 TestEnvironment.LogsEndpoint,
                 TestEnvironment.Credential,
-                InstrumentClientOptions(new LogsQueryClientOptions())
+                InstrumentClientOptions(new LogsQueryClientOptions()
+                {
+                    Diagnostics = { IsLoggingContentEnabled = true }
+                })
             ));
         }
 
@@ -444,6 +447,7 @@ public async Task ThrowsExceptionWhenBatchQueryNotFound()
 
             LogsBatchQuery batch = new LogsBatchQuery();
             batch.AddQuery(TestEnvironment.WorkspaceId, _logsTestData.TableAName, _logsTestData.DataTimeRange);
+
             var batchResult = await client.QueryBatchAsync(batch);
 
             var exception = Assert.Throws<ArgumentException>(() => batchResult.Value.GetResult("12345"));
@@ -502,16 +506,35 @@ public async Task CanQueryWithStatisticsBatch(bool include)
         }
 
         [RecordedTest]
-        public void CanSetServiceTimeout()
+        public async Task CanSetServiceTimeout()
         {
             var client = CreateClient();
 
-            var exception = Assert.ThrowsAsync<RequestFailedException>(async () => await client.QueryAsync(TestEnvironment.WorkspaceId, "range x from 1 to 100000000000 step 1 | count", _logsTestData.DataTimeRange, options: new LogsQueryOptions()
+            // Sometimes the service doesn't abort the query quickly enough
+            // and the request gets aborted instead
+            // Retry until we get the service to abort
+            while (true)
             {
-                ServerTimeout = TimeSpan.FromSeconds(1)
-            }));
-
-            Assert.AreEqual(504, exception.Status);
+                // Punch through caching
+                var cnt = 100000000000 + Recording.Random.Next(10000);
+                try
+                {
+                    await client.QueryAsync(TestEnvironment.WorkspaceId, $"range x from 1 to {cnt} step 1 | count", _logsTestData.DataTimeRange, options: new LogsQueryOptions()
+                    {
+                        ServerTimeout = TimeSpan.FromSeconds(1)
+                    });
+                }
+                catch (TaskCanceledException)
+                {
+                    // The client cancelled, retry.
+                    continue;
+                }
+                catch (RequestFailedException exception)
+                {
+                    Assert.AreEqual(504, exception.Status);
+                    return;
+                }
+            }
         }
 
         private record TestModel

sdk/monitor/Azure.Monitor.Query/tests/LogsClientSamples.cs

@@ -2,7 +2,10 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Threading.Tasks;
+using Azure.Core;
 using Azure.Core.TestFramework;
+using Azure.Monitor.Query.Models;
 using NUnit.Framework;
 
 namespace Azure.Monitor.Query.Tests
@@ -20,7 +23,7 @@ public void CanSetServiceTimeout_Mocked()
                 Assert.True(message.Request.Headers.TryGetValue("prefer", out preferHeader));
                 networkOverride = message.NetworkTimeout;
 
-                return new MockResponse(500);
+                return new MockResponse(403);
             });
 
             var client = new LogsQueryClient(new Uri("https://api.loganalytics.io"), new MockCredential(), new LogsQueryClientOptions()
@@ -34,7 +37,44 @@ public void CanSetServiceTimeout_Mocked()
             }));
 
             Assert.AreEqual("wait=600", preferHeader);
-            Assert.AreEqual(TimeSpan.FromMinutes(10), networkOverride);
+            // The network timeout is adjusted with 15 sec buffer
+            Assert.AreEqual(TimeSpan.FromMinutes(10).Add(TimeSpan.FromSeconds(15)), networkOverride);
+        }
+
+        [Test]
+        public async Task QueryBatchHandledInvalidResponse()
+        {
+            var badResponse = @"{
+    ""responses"": [
+        {
+            ""id"": ""0"",
+            ""status"": 200,
+            ""headers"": {
+                ""Age"": ""3"",
+                ""request-context"": ""appId=cid-v1:70941e4f-7e8f-40b7-b730-183893db0297""
+            },
+            ""body"": ""{\""tables\"":[{\""name\"":\""PrimaryResult\"",\""columns\"":[{\""name\"":\""TenantId\"",\""type\"":\""string\""},{\""name\"":\""SourceSystem\"",\""type\"":\""string\""},{\""name\"":\""MG\"",\""type\"":\""string\""},{\""name\"":\""ManagementGroupName\"",\""type\"":\""string\""},{\""name\"":\""TimeGenerated\"",\""type\"":\""datetime\""},{\""name\"":\""Computer\"",\""type\"":\""string\""},{\""name\"":\""RawData\"",\""type\"":\""string\""},{\""name\"":\""IntColumn_d\"",\""type\"":\""real\""},{\""name\"":\""StringColumn_s\"",\""type\"":\""string\""},{\""name\"":\""BoolColumn_b\"",\""type\"":\""bool\""},{\""name\"":\""FloatColumn_d\"",\""type\"":\""real\""},{\""name\"":\""Type\"",\""type\"":\""string\""},{\""name\"":\""_ResourceId\"",\""type\"":\""string\""}],\""rows\"":[[\""e7bf7412-576d-4978-b47c-2edf669e3e2a\"",\""RestAPI\"",\""\"",\""\"",\""2021-05-31T00:00:00Z\"",\""\"",\""\"",1,\""a\"",false,0,\""TableA1_151_CL\"",\""\""],[\""e7bf7412-576d-4978-b47c-2edf669e3e2a\"",\""RestAPI\"",\""\"",\""\"",\""2021-06-02T00:00:00Z\"",\""\"",\""\"",3,\""b\"",true,1.20000005,\""TableA1_151_CL\"",\""\""],[\""e7bf7412-576d-4978-b47c-2edf669e3e2a\"",\""RestAPI\"",\""\"",\""\"",\""2021-06-05T00:00:00Z\"",\""\"",\""\"",1,\""c\"",false,1.10000002,\""TableA1_151_CL\"",\""\""]]}]}""
+        }
+    ]
+}
+";
+            var mockTransport = MockTransport.FromMessageCallback(message =>
+            {
+                var mockResponse = new MockResponse(200);
+                mockResponse.SetContent(badResponse);
+                return mockResponse;
+            });
+
+            var client = new LogsQueryClient(new Uri("https://api.loganalytics.io"), new MockCredential(), new LogsQueryClientOptions()
+            {
+                Transport = mockTransport
+            });
+
+            LogsBatchQuery batch = new LogsBatchQuery();
+            batch.AddQuery("wid", "query", DateTimeRange.All);
+
+            LogsBatchQueryResult batchResult = await client.QueryBatchAsync(batch);
+            Assert.NotNull(batchResult.GetResult("0"));
         }
     }
 }

sdk/monitor/Azure.Monitor.Query/tests/LogsQueryClientClientLiveTests.cs

@@ -12,6 +12,7 @@ namespace Azure.Monitor.Query.Tests
     public class MetricsClientSamples: SamplesBase<MonitorQueryClientTestEnvironment>
     {
         [Test]
+        [Ignore("https://github.com/Azure/azure-sdk-for-net/issues/21657")]
         public async Task QueryMetrics()
         {
             #region Snippet:QueryMetrics

sdk/monitor/Azure.Monitor.Query/tests/LogsQueryClientClientTests.cs

@@ -31,7 +31,7 @@ private MetricsQueryClient CreateClient()
         [SetUp]
         public async Task SetUp()
         {
-            _testData = new MetricsTestData(this);
+            _testData = new MetricsTestData(TestEnvironment, Recording.UtcNow);
             await _testData.InitializeAsync();
         }