Add query by item test, fix query logic

Azure · Oct 10, 2022 · 5378b26 · 5378b26
1 parent e849e34
commit 5378b26
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/e2e/test/iothub/service/QueryClientE2ETests.cs b/e2e/test/iothub/service/QueryClientE2ETests.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net;
@@ -113,6 +114,42 @@ public async Task TwinQuery_CustomPaginationWorks()
             secondPageEnumerator.MoveNext().Should().BeFalse();
         }
 
+        [LoggedTestMethod]
+        [Timeout(TestTimeoutMilliseconds)]
+        public async Task TwinQuery_IterateByItemWorksAcrossPages()
+        {
+            // arrange
+
+            using var serviceClient = new IotHubServiceClient(TestConfiguration.IotHub.ConnectionString);
+            using TestDevice testDevice1 = await TestDevice.GetTestDeviceAsync(Logger, _idPrefix);
+            using TestDevice testDevice2 = await TestDevice.GetTestDeviceAsync(Logger, _idPrefix);
+            using TestDevice testDevice3 = await TestDevice.GetTestDeviceAsync(Logger, _idPrefix);
+
+            string queryText = $"select * from devices where deviceId = '{testDevice1.Id}' OR deviceId = '{testDevice2.Id}' OR deviceId = '{testDevice3.Id}'";
+            QueryOptions queryOptions = new QueryOptions
+            {
+                PageSize = 1
+            };
+
+            // act
+
+            await WaitForDevicesToBeQueryableAsync(serviceClient.Query, queryText, 3);
+
+            QueryResponse<Twin> twinQuery = await serviceClient.Query.CreateAsync<Twin>(queryText, queryOptions);
+
+            // assert
+            List<string> returnedTwinDeviceIds = new();
+            while (await twinQuery.MoveNextAsync())
+            {
+                Twin queriedTwin = twinQuery.Current;
+                returnedTwinDeviceIds.Add(queriedTwin.DeviceId);
+            }
+
+            List<string> expectedDeviceIds = new List<string>() { testDevice1.Id, testDevice2.Id, testDevice3.Id };
+            returnedTwinDeviceIds.Count.Should().Be(3);
+            returnedTwinDeviceIds.Should().Contain(expectedDeviceIds);
+        }
+
         [LoggedTestMethod]
         [Timeout(TestTimeoutMilliseconds)]
         public async Task JobQuery_QueryWorks()

diff --git a/iothub/service/src/Query/QueryResponse.cs b/iothub/service/src/Query/QueryResponse.cs
@@ -25,7 +25,7 @@ public class QueryResponse<T>
         private readonly JobType? _jobType;
         private readonly JobStatus? _jobStatus;
         private readonly int? _defaultPageSize;
-        private readonly IEnumerator<T> _items;
+        private IEnumerator<T> _items;
 
         internal QueryResponse(
             QueryClient client,
@@ -179,7 +179,7 @@ public async Task<bool> MoveNextAsync(QueryOptions queryOptions = default, Cance
                     .CreateAsync<T>(_originalQuery, queryOptionsClone, cancellationToken)
                     .ConfigureAwait(false);
                 CurrentPage = response.CurrentPage;
-                _items.Reset();
+                _items = response.CurrentPage.GetEnumerator();
                 _items.MoveNext();
                 Current = _items.Current;
                 ContinuationToken = response.ContinuationToken;