1 files changed, 95 insertions, 0 deletions
diff --git a/tensorflow/core/graph/edgeset_test.cc b/tensorflow/core/graph/edgeset_test.cc
new file mode 100644
index 0000000000..7909e8ea0a
--- /dev/null
+++ b/tensorflow/core/graph/edgeset_test.cc
@@ -0,0 +1,95 @@
+#include "tensorflow/core/graph/edgeset.h"
+
+#include "tensorflow/core/graph/graph.h"
+#include <gtest/gtest.h>
+
+namespace tensorflow {
+class EdgeSetTest : public ::testing::Test {
+ public:
+  EdgeSetTest() : edges_(nullptr), eset_(nullptr) {}
+
+  ~EdgeSetTest() override {
+    delete eset_;
+    delete[] edges_;
+  }
+
+  void MakeEdgeSet(int n) {
+    delete eset_;
+    delete[] edges_;
+    edges_ = new Edge[n];
+    eset_ = new EdgeSet;
+    model_.clear();
+    for (int i = 0; i < n; i++) {
+      eset_->insert(&edges_[i]);
+      model_.insert(&edges_[i]);
+    }
+  }
+
+  void CheckSame() {
+    EXPECT_EQ(model_.size(), eset_->size());
+    EXPECT_EQ(model_.empty(), eset_->empty());
+    std::vector<const Edge*> modelv(model_.begin(), model_.end());
+    std::vector<const Edge*> esetv(eset_->begin(), eset_->end());
+    std::sort(modelv.begin(), modelv.end());
+    std::sort(esetv.begin(), esetv.end());
+    EXPECT_EQ(modelv.size(), esetv.size());
+    for (size_t i = 0; i < modelv.size(); i++) {
+      EXPECT_EQ(modelv[i], esetv[i]) << i;
+    }
+  }
+
+  Edge nonexistent_;
+  Edge* edges_;
+  EdgeSet* eset_;
+  std::set<const Edge*> model_;
+};
+
+namespace {
+
+TEST_F(EdgeSetTest, Ops) {
+  for (int n : {0, 1, 2, 3, 4, 10}) {
+    MakeEdgeSet(n);
+    CheckSame();
+    EXPECT_EQ((n == 0), eset_->empty());
+    EXPECT_EQ(n, eset_->size());
+
+    eset_->clear();
+    model_.clear();
+    CheckSame();
+
+    eset_->insert(&edges_[0]);
+    model_.insert(&edges_[0]);
+    CheckSame();
+  }
+}
+
+// Try insert/erase of existing elements at different positions.
+TEST_F(EdgeSetTest, Exists) {
+  for (int n : {0, 1, 2, 3, 4, 10}) {
+    MakeEdgeSet(n);
+    for (int pos = 0; pos < n; pos++) {
+      MakeEdgeSet(n);
+      auto p = eset_->insert(&edges_[pos]);
+      EXPECT_FALSE(p.second);
+      EXPECT_EQ(&edges_[pos], *p.first);
+
+      EXPECT_EQ(1, eset_->erase(&edges_[pos]));
+      model_.erase(&edges_[pos]);
+      CheckSame();
+    }
+  }
+}
+
+// Try insert/erase of non-existent element.
+TEST_F(EdgeSetTest, DoesNotExist) {
+  for (int n : {0, 1, 2, 3, 4, 10}) {
+    MakeEdgeSet(n);
+    EXPECT_EQ(0, eset_->erase(&nonexistent_));
+    auto p = eset_->insert(&nonexistent_);
+    EXPECT_TRUE(p.second);
+    EXPECT_EQ(&nonexistent_, *p.first);
+  }
+}
+
+}  // namespace
+}  // namespace tensorflow