From 46c553f99ecb1616f242a888885a1e6eb61f865b Mon Sep 17 00:00:00 2001
From: George Pollard <gpollard@microsoft.com>
Date: Wed, 22 Jun 2022 04:07:38 +0000
Subject: [PATCH] More tests, improve SearchStatesQuery

---
 .../ApiService/onefuzzlib/NodeOperations.cs   |  4 +
 src/ApiService/Tests/Functions/NodeTests.cs   | 91 +++++++++++++++++++
 2 files changed, 95 insertions(+)

diff --git a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
index 9a0a426672c..1f12f83db4f 100644
--- a/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
+++ b/src/ApiService/ApiService/onefuzzlib/NodeOperations.cs
@@ -318,6 +318,10 @@ public static string SearchStatesQuery(
             queryParts.Add($"(pool_id eq '{poolId}')");
         }
 
+        if (poolName is not null) {
+            queryParts.Add($"(PartitionKey eq '{poolName}')");
+        }
+
         if (scaleSetId is not null) {
             queryParts.Add($"(scaleset_id eq '{scaleSetId}')");
         }
diff --git a/src/ApiService/Tests/Functions/NodeTests.cs b/src/ApiService/Tests/Functions/NodeTests.cs
index f959e81bcac..3b0a453a4a1 100644
--- a/src/ApiService/Tests/Functions/NodeTests.cs
+++ b/src/ApiService/Tests/Functions/NodeTests.cs
@@ -1,5 +1,6 @@
 
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Net;
 using Microsoft.OneFuzz.Service;
@@ -27,6 +28,7 @@ public NodeTestBase(ITestOutputHelper output, IStorage storage)
         : base(output, storage) { }
 
     private readonly Guid _machineId = Guid.NewGuid();
+    private readonly Guid _scalesetId = Guid.NewGuid();
     private readonly PoolName _poolName = PoolName.Parse($"pool-{Guid.NewGuid()}");
     private readonly string _version = Guid.NewGuid().ToString();
 
@@ -57,6 +59,95 @@ await Context.InsertAll(
         Assert.Equal(_version, deserialized.Version);
     }
 
+    [Fact]
+    public async Async.Task Search_MultipleNodes_CanFindNone() {
+        var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
+
+        var req = new NodeSearch();
+        var func = new NodeFunction(Logger, auth, Context);
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", req));
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+        Assert.Equal(0, result.Body.Length);
+    }
+
+    [Fact]
+    public async Async.Task Search_MultipleNodes_ByPoolName() {
+        await Context.InsertAll(
+            new Node(PoolName.Parse("otherPool"), Guid.NewGuid(), null, _version),
+            new Node(_poolName, Guid.NewGuid(), null, _version),
+            new Node(_poolName, Guid.NewGuid(), null, _version));
+
+        var req = new NodeSearch(PoolName: _poolName);
+
+        var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
+        var func = new NodeFunction(Logger, auth, Context);
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", req));
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        // make sure we got the data from the table
+        var deserialized = BodyAs<NodeSearchResult[]>(result);
+        Assert.Equal(2, deserialized.Length);
+    }
+
+    [Fact]
+    public async Async.Task Search_MultipleNodes_ByScalesetId() {
+        await Context.InsertAll(
+            new Node(_poolName, Guid.NewGuid(), null, _version, ScalesetId: _scalesetId),
+            new Node(_poolName, Guid.NewGuid(), null, _version, ScalesetId: _scalesetId),
+            new Node(_poolName, Guid.NewGuid(), null, _version));
+
+        var req = new NodeSearch(ScalesetId: _scalesetId);
+
+        var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
+        var func = new NodeFunction(Logger, auth, Context);
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", req));
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        // make sure we got the data from the table
+        var deserialized = BodyAs<NodeSearchResult[]>(result);
+        Assert.Equal(2, deserialized.Length);
+    }
+
+
+    [Fact]
+    public async Async.Task Search_MultipleNodes_ByState() {
+        await Context.InsertAll(
+            new Node(_poolName, Guid.NewGuid(), null, _version, State: NodeState.Busy),
+            new Node(_poolName, Guid.NewGuid(), null, _version, State: NodeState.Busy),
+            new Node(_poolName, Guid.NewGuid(), null, _version));
+
+        var req = new NodeSearch(State: new List<NodeState> { NodeState.Busy });
+
+        var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
+        var func = new NodeFunction(Logger, auth, Context);
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", req));
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        // make sure we got the data from the table
+        var deserialized = BodyAs<NodeSearchResult[]>(result);
+        Assert.Equal(2, deserialized.Length);
+    }
+
+    [Fact]
+    public async Async.Task Search_MultipleNodes_ByMultipleStates() {
+        await Context.InsertAll(
+            new Node(_poolName, Guid.NewGuid(), null, _version, State: NodeState.Free),
+            new Node(_poolName, Guid.NewGuid(), null, _version, State: NodeState.Busy),
+            new Node(_poolName, Guid.NewGuid(), null, _version, State: NodeState.Busy),
+            new Node(_poolName, Guid.NewGuid(), null, _version));
+
+        var req = new NodeSearch(State: new List<NodeState> { NodeState.Free, NodeState.Busy });
+
+        var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
+        var func = new NodeFunction(Logger, auth, Context);
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", req));
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        // make sure we got the data from the table
+        var deserialized = BodyAs<NodeSearchResult[]>(result);
+        Assert.Equal(3, deserialized.Length);
+    }
+
     [Theory]
     [InlineData("PATCH")]
     [InlineData("POST")]