ADT: Add a new Query Test that does a COUNT (Azure#16586)

Author: Basel Rustum

sdk/digitaltwins/Azure.DigitalTwins.Core/tests/QueryTests.cs

@@ -6,13 +6,17 @@
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using Azure.Core.TestFramework;
 using FluentAssertions;
 using NUnit.Framework;
 
 namespace Azure.DigitalTwins.Core.Tests
 {
     public class QueryTests : E2eTestBase
     {
+        private static readonly int s_retryCount = 10;
+        private static readonly TimeSpan s_retryDelay = TimeSpan.FromSeconds(2);
+
         public QueryTests(bool isAsync)
             : base(isAsync)
         {
@@ -37,19 +41,38 @@ public async Task Query_ValidQuery_Success()
                 // Create a room twin, with property "IsOccupied": true
                 string roomTwinId = await GetUniqueTwinIdAsync(client, TestAssetDefaults.RoomTwinIdPrefix).ConfigureAwait(false);
                 BasicDigitalTwin roomTwin = TestAssetsHelper.GetRoomTwinPayload(roomModelId);
-                await client.CreateOrReplaceDigitalTwinAsync<BasicDigitalTwin>(roomTwinId, roomTwin).ConfigureAwait(false);
+                await client.CreateOrReplaceDigitalTwinAsync(roomTwinId, roomTwin).ConfigureAwait(false);
 
-                string queryString = "SELECT * FROM digitaltwins where IsOccupied = true";
+                // Construct a query string to find the twins with the EXACT model id and provided version. If EXACT is not specified, the query
+                // call will get all twins with the same model id but that implement any version higher than the provided version
+                string queryString = $"SELECT * FROM digitaltwins WHERE IS_OF_MODEL('{roomModelId}', EXACT) AND IsOccupied = true";
 
                 // act
-                AsyncPageable<JsonElement> asyncPageableResponse = client.QueryAsync<JsonElement>(queryString);
+                AsyncPageable<BasicDigitalTwin> asyncPageableResponse = client.QueryAsync<BasicDigitalTwin>(queryString);
 
                 // assert
-                await foreach (JsonElement response in asyncPageableResponse)
+
+                // It takes a few seconds for the service to be able to fetch digital twins through queries after being created. Hence, adding the retry logic
+                var digitalTwinFound = false;
+                await TestRetryHelper.RetryAsync<AsyncPageable<BasicDigitalTwin>>(async () =>
                 {
-                    JsonElement isOccupied = response.GetProperty("IsOccupied");
-                    isOccupied.GetRawText().Should().Be("true");
-                }
+                    await foreach (BasicDigitalTwin response in asyncPageableResponse)
+                    {
+                        digitalTwinFound = true;
+                        bool isOccupied = ((JsonElement)response.Contents["IsOccupied"]).GetBoolean();
+                        isOccupied.Should().BeTrue();
+                        break;
+                    }
+
+                    if (!digitalTwinFound)
+                    {
+                        throw new Exception($"Digital twin based on model Id {roomModelId} not found");
+                    }
+
+                    return null;
+                }, s_retryCount, s_retryDelay);
+
+                digitalTwinFound.Should().BeTrue();
             }
             catch (Exception ex)
             {
@@ -154,5 +177,78 @@ public async Task Query_PaginationWorks()
                 }
             }
         }
+
+        [Test]
+        public async Task Query_GetTwinCount()
+        {
+            DigitalTwinsClient client = GetClient();
+
+            string floorModelId = await GetUniqueModelIdAsync(client, TestAssetDefaults.FloorModelIdPrefix).ConfigureAwait(false);
+            string roomModelId = await GetUniqueModelIdAsync(client, TestAssetDefaults.RoomModelIdPrefix).ConfigureAwait(false);
+
+            try
+            {
+                // arrange
+
+                // Create room model
+                string roomModel = TestAssetsHelper.GetRoomModelPayload(roomModelId, floorModelId);
+                await CreateAndListModelsAsync(client, new List<string> { roomModel }).ConfigureAwait(false);
+
+                // Create a room twin, with property "IsOccupied": true
+                string roomTwinId = await GetUniqueTwinIdAsync(client, TestAssetDefaults.RoomTwinIdPrefix).ConfigureAwait(false);
+                BasicDigitalTwin roomTwin = TestAssetsHelper.GetRoomTwinPayload(roomModelId);
+                await client.CreateOrReplaceDigitalTwinAsync(roomTwinId, roomTwin).ConfigureAwait(false);
+
+                // Construct a query string to find the twins with the EXACT model id and provided version. If EXACT is not specified, the query
+                // call will get all twins with the same model id but that implement any version higher than the provided version
+                string queryString = $"SELECT COUNT() FROM digitaltwins WHERE IS_OF_MODEL('{roomModelId}', EXACT) AND IsOccupied = true";
+
+                // act
+                AsyncPageable<JsonElement> asyncPageableResponse = client.QueryAsync<JsonElement>(queryString);
+
+                // assert
+
+                // It takes a few seconds for the service to be able to fetch digital twins through queries after being created. Hence, adding the retry logic
+                var currentCount = 0;
+                await TestRetryHelper.RetryAsync<AsyncPageable<JsonElement>>(async () =>
+                {
+                    await foreach (JsonElement response in asyncPageableResponse)
+                    {
+                        string currentCountStr = response.GetRawText();
+                        IDictionary<string, int> currentCountDictionary = JsonSerializer.Deserialize<IDictionary<string, int>>(currentCountStr);
+                        currentCountDictionary.ContainsKey("COUNT").Should().BeTrue();
+                        currentCount = currentCountDictionary["COUNT"];
+                    }
+
+                    if (currentCount == 0)
+                    {
+                        throw new Exception($"Digital twin based on model Id {roomModelId} not found");
+                    }
+
+                    return null;
+                }, s_retryCount, s_retryDelay);
+
+                currentCount.Should().Be(1);
+            }
+            catch (Exception ex)
+            {
+                Assert.Fail($"Failure in executing a step in the test case: {ex.Message}.");
+            }
+            finally
+            {
+                // clean up
+                try
+                {
+                    if (!string.IsNullOrWhiteSpace(roomModelId))
+                    {
+                        await client.DeleteModelAsync(roomModelId).ConfigureAwait(false);
+                    }
+                }
+                catch (Exception ex)
+                {
+                    Assert.Fail($"Test clean up failed: {ex.Message}");
+                }
+            }
+        }
     }
 }