From a584f19afb9dd05b59052b21d79eae5252c783e2 Mon Sep 17 00:00:00 2001
From: Arkadiusz Szczepkowicz <arek.szczepkowicz@ng-analytics.com>
Date: Mon, 22 May 2023 18:33:35 +0200
Subject: [PATCH] #2074: Add tests for detecting OfflineLB in LB data file

---
 src/vt/vrt/collection/balance/read_lb.cc | 15 ++++++---------
 src/vt/vrt/collection/balance/read_lb.h  |  3 ++-
 tests/unit/lb/test_lb_reader.nompi.cc    | 18 ++++++++++++++++++
 tests/unit/lb/test_offlinelb.cc          |  1 -
 4 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/vt/vrt/collection/balance/read_lb.cc b/src/vt/vrt/collection/balance/read_lb.cc
index a32e99e45b..d85edf680a 100644
--- a/src/vt/vrt/collection/balance/read_lb.cc
+++ b/src/vt/vrt/collection/balance/read_lb.cc
@@ -61,6 +61,7 @@ namespace vt { namespace vrt { namespace collection { namespace balance {
 /*static*/ typename ReadLBConfig::ConfigMapType ReadLBConfig::config_exact_ = {};
 /*static*/ std::vector<ConfigIndex> ReadLBConfig::config_prec_ = {};
 /*static*/ bool ReadLBConfig::read_complete_ = false;
+/*static*/ bool ReadLBConfig::has_offline_lb_ = false;
 
 /*static*/ bool ReadLBConfig::openConfig(std::string const& filename) {
   // No-op if no file specified. Can't be used to clear.
@@ -102,15 +103,6 @@ namespace vt { namespace vrt { namespace collection { namespace balance {
   }
 }
 
-/*static*/ bool ReadLBConfig::hasOfflineLB() {
-  for(auto&& ele : config_exact_) {
-    if(getLB(ele.first) == LBType::OfflineLB) {
-      return true;
-    }
-  }
-  return false;
-}
-
 /*static*/ ConfigEntry* ReadLBConfig::entry(ConfigIndex const& idx) {
   // First, search the exact iter config for this iteration: it has the highest
   // precedence
@@ -240,6 +232,10 @@ int eatWhitespace(std::ifstream& file) {
       vtAbort(err_msg);
     }
 
+    if (lb_name == get_lb_names()[LBType::OfflineLB]) {
+      has_offline_lb_ = true;
+    }
+
     map->emplace(
       std::piecewise_construct,
       std::forward_as_tuple(mod),
@@ -252,6 +248,7 @@ int eatWhitespace(std::ifstream& file) {
 
 /*static*/ void ReadLBConfig::clear() {
   read_complete_ = false;
+  has_offline_lb_ = false;
   open_filename_ = "";
   config_mod_.clear();
   config_exact_.clear();
diff --git a/src/vt/vrt/collection/balance/read_lb.h b/src/vt/vrt/collection/balance/read_lb.h
index 60f4c88dd0..03d409307f 100644
--- a/src/vt/vrt/collection/balance/read_lb.h
+++ b/src/vt/vrt/collection/balance/read_lb.h
@@ -197,7 +197,7 @@ struct ReadLBConfig {
   static ConfigIndex numEntries() { return config_mod_.size() + config_exact_.size(); }
   static ConfigEntry* entry(ConfigIndex const& idx);
   static LBType getLB(ConfigIndex const& idx);
-  static bool hasOfflineLB();
+  static bool hasOfflineLB() { return has_offline_lb_; };
   static ConfigMapType getModEntries() { return config_mod_; };
   static ConfigMapType getExactEntries() {return config_exact_; };
   static ParamMapType parseParams(std::vector<std::string> params);
@@ -209,6 +209,7 @@ struct ReadLBConfig {
   static void readFile(std::string const& filename);
 
   static bool read_complete_;
+  static bool has_offline_lb_;
   static std::string open_filename_;
   static ConfigMapType config_mod_;
   static ConfigMapType config_exact_;
diff --git a/tests/unit/lb/test_lb_reader.nompi.cc b/tests/unit/lb/test_lb_reader.nompi.cc
index 3451bab7a7..06c74b125b 100644
--- a/tests/unit/lb/test_lb_reader.nompi.cc
+++ b/tests/unit/lb/test_lb_reader.nompi.cc
@@ -66,6 +66,7 @@ TEST_F(TestLBReader, test_lb_read_1) {
   Config::clear();
   Config::openConfig(file_name);
 
+  EXPECT_EQ(Config::hasOfflineLB(), false);
   EXPECT_EQ(Config::numEntries(), 3);
   EXPECT_EQ(Config::getExactEntries().size(), 2);
   EXPECT_EQ(Config::getModEntries().size(), 1);
@@ -120,6 +121,7 @@ TEST_F(TestLBReader, test_lb_read_2) {
   Config::clear();
   Config::openConfig(file_name);
 
+  EXPECT_EQ(Config::hasOfflineLB(), false);
   EXPECT_EQ(Config::numEntries(), 5);
   for (ConfigIdx i = 0; i < 121; i++) {
     auto entry = Config::entry(i);
@@ -195,4 +197,20 @@ TEST_F(TestLBReader, test_lb_read_2) {
   EXPECT_EQ(Config::toString(), expected_config);
 }
 
+TEST_F(TestLBReader, test_lb_read_3_with_offline_lb) {
+  std::string file_name = "test_lb_read_3.txt";
+  std::ofstream out(file_name);
+  out << ""
+    "0 NoLB\n"
+    "1 OfflineLB\n"
+    "%10 OfflineLB\n";
+  out.close();
+
+  using Config = vt::vrt::collection::balance::ReadLBConfig;
+
+  Config::clear();
+  Config::openConfig(file_name);
+  EXPECT_EQ(Config::hasOfflineLB(), true);
+}
+
 }}} // end namespace vt::tests::unit
diff --git a/tests/unit/lb/test_offlinelb.cc b/tests/unit/lb/test_offlinelb.cc
index eaf1bf6750..f46f096fcc 100644
--- a/tests/unit/lb/test_offlinelb.cc
+++ b/tests/unit/lb/test_offlinelb.cc
@@ -76,7 +76,6 @@ struct SimCol : vt::Collection<SimCol, vt::Index1D> {
   }
 
   void sparseHandler(Msg* m){
-    // TODO
   }
 };