random_selector: Add method for removing entries

BUG=chromium:371228
TEST=unit tests pass

Change-Id: I9a0d7d6daa57a2b3215f3a3269afd2c5f9f736ea
Signed-off-by: Simon Que <sque@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/199662
Reviewed-by: Emma Rapati <rapati@chromium.org>
diff --git a/src/random_selector.cc b/src/random_selector.cc
index b81c7e4..62fc6bd 100644
--- a/src/random_selector.cc
+++ b/src/random_selector.cc
@@ -67,6 +67,13 @@
   GetKeyOf(random, next);
 }
 
+void RandomSelector::Remove(const std::string& key) {
+  std::map<std::string, float>::iterator iter = odds_.find(key);
+  if (iter != odds_.end()) {
+    odds_.erase(iter);
+  }
+}
+
 float RandomSelector::GetFloatBetween(float min, float max) {
   CHECK_GT(max, min);
   float random = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
diff --git a/src/random_selector.h b/src/random_selector.h
index d1de3bb..989d65d 100644
--- a/src/random_selector.h
+++ b/src/random_selector.h
@@ -43,6 +43,15 @@
   // Get the next randomly picked string in |next|.
   void GetNext(std::string* next);
 
+  // Removes an entry from the current odds container. If the entry is not found
+  // in |odds_|, this function has no effect.
+  void Remove(const std::string& key);
+
+  // Returns the number of string entries.
+  size_t GetNumStrings() const {
+    return odds_.size();
+  }
+
  private:
   // Get a floating point number between |min| and |max|.
   virtual float GetFloatBetween(float min, float max);
diff --git a/src/random_selector_test.cc b/src/random_selector_test.cc
index 32b2127..2cad7e0 100644
--- a/src/random_selector_test.cc
+++ b/src/random_selector_test.cc
@@ -111,3 +111,20 @@
   GenerateResults(odds.size(), &random_selector, &results);
   CheckResultsAgainstOdds(odds, results);
 }
+
+// Ensure RandomSelector is able to delete odds that it has previously stored.
+TEST(RandomSelector, RemoveEntriesTest) {
+  RandomSelector random_selector;
+  random_selector.SetOddsFromFile(std::string(kOddsFilename));
+  std::string key;
+
+  // Get a key and verify.
+  random_selector.GetNext(&key);
+  EXPECT_FALSE(key.empty());
+  size_t old_size = random_selector.GetNumStrings();
+  EXPECT_GT(old_size, 0);
+
+  // Remove the key.
+  random_selector.Remove(key);
+  EXPECT_EQ(old_size - 1, random_selector.GetNumStrings());
+}