From ccf99e8b411657c86939be9d312661cbcf854ab8 Mon Sep 17 00:00:00 2001
From: Carlos <carlos@carlossegarra.com>
Date: Fri, 4 Aug 2023 12:27:02 +0100
Subject: [PATCH] Add method to get unique hosts from scheduling decision
 (#346)

decisions: add a couple of utility methods
---
 .../faabric/batch-scheduler/SchedulingDecision.h    |  4 ++++
 src/batch-scheduler/SchedulingDecision.cpp          |  5 +++++
 .../batch-scheduler/test_scheduling_decisions.cpp   | 13 ++++++++++++-
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/include/faabric/batch-scheduler/SchedulingDecision.h b/include/faabric/batch-scheduler/SchedulingDecision.h
index 5b8467b26..d28475a10 100644
--- a/include/faabric/batch-scheduler/SchedulingDecision.h
+++ b/include/faabric/batch-scheduler/SchedulingDecision.h
@@ -63,6 +63,8 @@ class SchedulingDecision
 
     SchedulingDecision(uint32_t appIdIn, int32_t groupIdIn);
 
+    bool operator==(const SchedulingDecision& rhs) const = default;
+
     uint32_t appId = 0;
 
     int32_t groupId = 0;
@@ -96,6 +98,8 @@ class SchedulingDecision
                     int32_t messageId,
                     int32_t appIdx,
                     int32_t groupIdx);
+
+    std::set<std::string> uniqueHosts();
 };
 
 }
diff --git a/src/batch-scheduler/SchedulingDecision.cpp b/src/batch-scheduler/SchedulingDecision.cpp
index 41325c963..d98e41b1e 100644
--- a/src/batch-scheduler/SchedulingDecision.cpp
+++ b/src/batch-scheduler/SchedulingDecision.cpp
@@ -52,4 +52,9 @@ SchedulingDecision SchedulingDecision::fromPointToPointMappings(
 
     return decision;
 }
+
+std::set<std::string> SchedulingDecision::uniqueHosts()
+{
+    return std::set<std::string>(hosts.begin(), hosts.end());
+}
 }
diff --git a/tests/test/batch-scheduler/test_scheduling_decisions.cpp b/tests/test/batch-scheduler/test_scheduling_decisions.cpp
index 56e6816c1..fe8ddb5c6 100644
--- a/tests/test/batch-scheduler/test_scheduling_decisions.cpp
+++ b/tests/test/batch-scheduler/test_scheduling_decisions.cpp
@@ -20,9 +20,10 @@ TEST_CASE_METHOD(ConfFixture, "Test building scheduling decisions", "[util]")
     std::string hostC = "hostC";
 
     std::string thisHost = conf.endpointHost;
+    std::set<std::string> expectedUniqueHosts;
 
     bool expectSingleHost = false;
-    SECTION("Multi-host") {}
+    SECTION("Multi-host") { expectedUniqueHosts = { hostA, hostB, hostC }; }
 
     SECTION("Only remote hosts")
     {
@@ -30,6 +31,7 @@ TEST_CASE_METHOD(ConfFixture, "Test building scheduling decisions", "[util]")
         hostC = "hostA";
 
         expectSingleHost = false;
+        expectedUniqueHosts = { "hostA" };
     }
 
     SECTION("All this host")
@@ -39,6 +41,7 @@ TEST_CASE_METHOD(ConfFixture, "Test building scheduling decisions", "[util]")
         hostC = thisHost;
 
         expectSingleHost = true;
+        expectedUniqueHosts = { thisHost };
     }
 
     SECTION("All this host single host optimisations off")
@@ -50,6 +53,7 @@ TEST_CASE_METHOD(ConfFixture, "Test building scheduling decisions", "[util]")
         hostC = thisHost;
 
         expectSingleHost = false;
+        expectedUniqueHosts = { thisHost };
     }
 
     auto req = faabric::util::batchExecFactory("foo", "bar", 3);
@@ -75,8 +79,15 @@ TEST_CASE_METHOD(ConfFixture, "Test building scheduling decisions", "[util]")
     REQUIRE(decision.nFunctions == 3);
     REQUIRE(decision.messageIds == expectedMsgIds);
     REQUIRE(decision.hosts == expectedHosts);
+    REQUIRE(decision.uniqueHosts() == expectedUniqueHosts);
     REQUIRE(decision.appIdxs == expectedAppIdxs);
     REQUIRE(decision.isSingleHost() == expectSingleHost);
+
+    // We can compare decisions with the == operator
+    auto newDecision = decision;
+    REQUIRE(newDecision == decision);
+    newDecision.groupId = 1338;
+    REQUIRE(newDecision != decision);
 }
 
 TEST_CASE("Test converting point-to-point mappings to scheduling decisions",