#2229: Add tests for reading and writing user attributes in LBDataHolder

DARMA-tasking · Feb 20, 2024 · 41ac48e · 41ac48e
1 parent e730ffd
commit 41ac48e
Showing 1 changed file with 70 additions and 5 deletions.
diff --git a/tests/unit/collection/test_lb_data_holder.cc b/tests/unit/collection/test_lb_data_holder.cc
@@ -58,7 +58,7 @@ namespace vt { namespace tests { namespace unit { namespace lb {
 
 static constexpr int const num_phases = 10;
 
-struct TestOptionalPhasesMetadata : TestParallelHarnessParam<bool> { };
+struct TestLBDataHolder : TestParallelHarnessParam<bool> { };
 
 void addPhasesDataToJson(nlohmann::json& json, PhaseType amountOfPhasesToAdd, std::vector<PhaseType> toSkip) {
   using ElementIDStruct = vt::vrt::collection::balance::ElementIDStruct;
@@ -90,7 +90,7 @@ void addPhasesDataToJson(nlohmann::json& json, PhaseType amountOfPhasesToAdd, st
   json["phases"] = phases;
 }
 
-TEST_F(TestOptionalPhasesMetadata, test_no_metadata) {
+TEST_F(TestLBDataHolder, test_no_metadata) {
   using LBDataHolder = vt::vrt::collection::balance::LBDataHolder;
 
   nlohmann::json json;
@@ -117,7 +117,7 @@ TEST_F(TestOptionalPhasesMetadata, test_no_metadata) {
   EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges);
 }
 
-TEST_F(TestOptionalPhasesMetadata, test_no_lb_phases_metadata) {
+TEST_F(TestLBDataHolder, test_no_lb_phases_metadata) {
   using LBDataHolder = vt::vrt::collection::balance::LBDataHolder;
 
   nlohmann::json json;
@@ -146,7 +146,7 @@ TEST_F(TestOptionalPhasesMetadata, test_no_lb_phases_metadata) {
   EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges);
 }
 
-TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_empty) {
+TEST_F(TestLBDataHolder, test_lb_phases_metadata_empty) {
   using LBDataHolder = vt::vrt::collection::balance::LBDataHolder;
 
   nlohmann::json metadata, phasesMetadata, json;
@@ -182,7 +182,7 @@ TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_empty) {
   EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges);
 }
 
-TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_filled) {
+TEST_F(TestLBDataHolder, test_lb_phases_metadata_filled) {
   using LBDataHolder = vt::vrt::collection::balance::LBDataHolder;
 
   nlohmann::json metadata, phasesMetadata, json;
@@ -218,6 +218,71 @@ TEST_F(TestOptionalPhasesMetadata, test_lb_phases_metadata_filled) {
   EXPECT_EQ((*outJsonPtr)["identical_to_previous"]["range"], expectedIdenticalRanges);
 }
 
+TEST_F(TestLBDataHolder, test_lb_rank_attributes) {
+  using LBDataHolder = vt::vrt::collection::balance::LBDataHolder;
+
+  nlohmann::json json = R"(
+    {
+      "type": "LBDatafile",
+      "metadata" : {
+        "attributes": {
+            "some_val": 123
+        }
+      },
+      "phases" : []
+    }
+  )"_json;
+
+  LBDataHolder testObj(json);
+  EXPECT_TRUE(nullptr != testObj.rank_attributes_);
+  EXPECT_EQ(123, (*testObj.rank_attributes_)["some_val"]);
+
+  auto outJsonPtr = testObj.metadataToJson();
+  ASSERT_TRUE(outJsonPtr != nullptr);
+  EXPECT_EQ(123, (*outJsonPtr)["attributes"]["some_val"]);
+}
+
+TEST_F(TestLBDataHolder, test_lb_entity_attributes) {
+  using LBDataHolder = vt::vrt::collection::balance::LBDataHolder;
+
+  nlohmann::json json = R"(
+    {
+      "type": "LBDatafile",
+      "phases": [
+        {
+          "id": 0,
+          "tasks": [
+            {
+              "entity": {
+                "home": 0,
+                "id": 524291,
+                "type": "object",
+                "migratable": true,
+                "attributes": {
+                  "some_val": 123
+                }
+              },
+              "node": 0,
+              "resource": "cpu",
+              "time": 3.0
+            }
+          ]
+        }
+      ]
+    }
+  )"_json;
+  auto id = vt::vrt::collection::balance::ElementIDStruct{524291, 0};
+
+  LBDataHolder testObj(json);
+  EXPECT_TRUE(testObj.node_user_attributes_.find(0) != testObj.node_user_attributes_.end());
+  EXPECT_TRUE(testObj.node_user_attributes_[0].find(id) != testObj.node_user_attributes_[0].end());
+  EXPECT_EQ(123, (*testObj.node_user_attributes_[0][id])["some_val"]);
+
+  auto outJsonPtr = testObj.toJson(0);
+  ASSERT_TRUE(outJsonPtr != nullptr);
+  EXPECT_EQ(123, (*outJsonPtr)["tasks"][0]["entity"]["attributes"]["some_val"]);
+}
+
 }}}} // end namespace vt::tests::unit::lb
 
 #endif /*vt_check_enabled(lblite)*/