Add tests for new methods in clean_relation.

probcomp · Oct 4, 2024 · 3110c4f · 3110c4f
1 parent 9ab0e4f
commit 3110c4f
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/cxx/clean_relation.hh b/cxx/clean_relation.hh
@@ -117,7 +117,6 @@ class CleanRelation : public Relation<T> {
     return value;
   }
 
-  // TODO: add a test.
   void cleanup_data(const T_items& items) {
     for (int i = 0; i < std::ssize(domains); ++i) {
       const std::string& name = domains[i]->name;
@@ -131,7 +130,6 @@ class CleanRelation : public Relation<T> {
     data.erase(items);
   }
 
-  // TODO: add a test.
   void cleanup_clusters() {
     for (auto it = clusters.cbegin(); it != clusters.cend();) {
       if (it->second->N == 0) {

diff --git a/cxx/clean_relation_test.cc b/cxx/clean_relation_test.cc
@@ -210,3 +210,55 @@ BOOST_AUTO_TEST_CASE(test_sample_and_incorporate) {
 
   BOOST_TEST(R1.data.size() == 3);
 }
+
+BOOST_AUTO_TEST_CASE(test_cleanup_data) {
+  std::mt19937 prng;
+  Domain D1("D1");
+  Domain D2("D2");
+  DistributionSpec spec("normal");
+  CleanRelation<double> R1("R1", spec, {&D1, &D2});
+  R1.incorporate(&prng, {0, 1}, 3.);
+  R1.incorporate(&prng, {0, 2}, 2.8);
+  R1.incorporate(&prng, {5, 2}, 2.3);
+
+  R1.unincorporate_from_cluster({5, 2});
+  BOOST_TEST(R1.data.contains({5, 2}));
+  BOOST_TEST(R1.data_r.at("D1").contains(5));
+  BOOST_TEST(R1.data_r.at("D2").contains(2));
+  R1.cleanup_data({5, 2});
+  BOOST_TEST(!R1.data.contains({5, 2}));
+  BOOST_TEST(R1.data.contains({0, 2}));
+  BOOST_TEST(!R1.data_r.at("D1").contains(5));
+  BOOST_TEST(R1.data_r.at("D1").contains(0));
+  BOOST_TEST(R1.data_r.at("D2").contains(2));
+}
+
+BOOST_AUTO_TEST_CASE(test_cleanup_clusters) {
+  std::mt19937 prng;
+  Domain D1("D1");
+  Domain D2("D2");
+  DistributionSpec spec("normal");
+  CleanRelation<double> R1("R1", spec, {&D1, &D2});
+  D1.incorporate(&prng, 0, 0);
+  D1.incorporate(&prng, 5, 0);
+  D2.incorporate(&prng, 1, 0);
+  D2.incorporate(&prng, 2, 1);
+  D2.incorporate(&prng, 3, 4);
+
+  R1.incorporate_items(&prng, {0, 1});
+  R1.incorporate_items(&prng, {0, 2});
+  R1.incorporate_items(&prng, {5, 2});
+  R1.incorporate_to_cluster({0, 1}, 3.);
+  R1.incorporate_to_cluster({0, 2}, 2.8);
+  BOOST_TEST(R1.clusters_contains({0, 1}));
+  BOOST_TEST(R1.clusters_contains({0, 2}));
+  BOOST_TEST(R1.clusters_contains({5, 2}));
+
+  R1.data[{0, 2}] = 2.8;
+  R1.unincorporate_from_cluster({0, 2});
+  R1.cleanup_clusters();
+  BOOST_TEST(R1.clusters_contains({0, 1}));
+  BOOST_TEST(R1.clusters_contains({5, 1}));
+  BOOST_TEST(!R1.clusters_contains({0, 2}));
+  BOOST_TEST(!R1.clusters_contains({5, 2}));
+}