4 files changed, 66 insertions, 3 deletions

diff --git a/components/rappor/rappor_service.h b/components/rappor/rappor_service.h
index 4f26214..d50b91e 100644
--- a/components/rappor/rappor_service.h
+++ b/components/rappor/rappor_service.h
@@ -65,9 +65,9 @@ class RapporService {
 
   // Records a sample of the rappor metric specified by |metric_name|.
   // Creates and initializes the metric, if it doesn't yet exist.
-  void RecordSample(const std::string& metric_name,
-                    RapporType type,
-                    const std::string& sample);
+  virtual void RecordSample(const std::string& metric_name,
+                            RapporType type,
+                            const std::string& sample);
 
   // Registers the names of all of the preferences used by RapporService in the
   // provided PrefRegistry. This should be called before calling Start().
diff --git a/components/rappor/rappor_service_unittest.cc b/components/rappor/rappor_service_unittest.cc
index 7fedd6e..eefdb52e 100644
--- a/components/rappor/rappor_service_unittest.cc
+++ b/components/rappor/rappor_service_unittest.cc
@@ -65,6 +65,24 @@ TEST(RapporServiceTest, RecordingLevel) {
   EXPECT_EQ(0, reports.report_size());
 }
 
+// Check that GetRecordedSampleForMetric works as expected.
+TEST(RapporServiceTest, GetRecordedSampleForMetric) {
+  TestRapporService rappor_service;
+
+  // Multiple samples for the same metric; only the latest is remembered.
+  rappor_service.RecordSample("MyMetric", ETLD_PLUS_ONE_RAPPOR_TYPE, "foo");
+  rappor_service.RecordSample("MyMetric", ETLD_PLUS_ONE_RAPPOR_TYPE, "bar");
+
+  std::string sample;
+  RapporType type;
+  EXPECT_FALSE(
+      rappor_service.GetRecordedSampleForMetric("WrongMetric", &sample, &type));
+  EXPECT_TRUE(
+      rappor_service.GetRecordedSampleForMetric("MyMetric", &sample, &type));
+  EXPECT_EQ("bar", sample);
+  EXPECT_EQ(ETLD_PLUS_ONE_RAPPOR_TYPE, type);
+}
+
 // Check that the incognito is respected.
 TEST(RapporServiceTest, Incognito) {
   TestRapporService rappor_service;
diff --git a/components/rappor/test_rappor_service.cc b/components/rappor/test_rappor_service.cc
index f355cb6..aa54a94 100644
--- a/components/rappor/test_rappor_service.cc
+++ b/components/rappor/test_rappor_service.cc
@@ -34,6 +34,18 @@ TestRapporService::TestRapporService()
 
 TestRapporService::~TestRapporService() {}
 
+void TestRapporService::RecordSample(const std::string& metric_name,
+                                     RapporType type,
+                                     const std::string& sample) {
+  // Save the recorded sample to the local structure.
+  RapporSample rappor_sample;
+  rappor_sample.type = type;
+  rappor_sample.value = sample;
+  samples_[metric_name] = rappor_sample;
+  // Original version is still called.
+  RapporService::RecordSample(metric_name, type, sample);
+}
+
 int TestRapporService::GetReportsCount() {
   RapporReports reports;
   ExportMetrics(&reports);
@@ -44,6 +56,18 @@ void TestRapporService::GetReports(RapporReports* reports) {
   ExportMetrics(reports);
 }
 
+bool TestRapporService::GetRecordedSampleForMetric(
+    const std::string& metric_name,
+    std::string* sample,
+    RapporType* type) {
+  SamplesMap::iterator it = samples_.find(metric_name);
+  if (it == samples_.end())
+    return false;
+  *sample = it->second.value;
+  *type = it->second.type;
+  return true;
+}
+
 // Cancel the next call to OnLogInterval.
 void TestRapporService::CancelNextLogRotation() {
   next_rotation_ = base::TimeDelta();
diff --git a/components/rappor/test_rappor_service.h b/components/rappor/test_rappor_service.h
index ee43742..ead80a8 100644
--- a/components/rappor/test_rappor_service.h
+++ b/components/rappor/test_rappor_service.h
@@ -5,6 +5,7 @@
 #ifndef COMPONENTS_RAPPOR_TEST_RAPPOR_SERVICE_H_
 #define COMPONENTS_RAPPOR_TEST_RAPPOR_SERVICE_H_
 
+#include <map>
 #include <string>
 
 #include "base/prefs/testing_pref_service.h"
@@ -22,6 +23,11 @@ class TestRapporService : public RapporService {
 
   ~TestRapporService() override;
 
+  // Intercepts the sample being recorded and saves it in a test structure.
+  void RecordSample(const std::string& metric_name,
+                    RapporType type,
+                    const std::string& sample) override;
+
   // Gets the number of reports that would be uploaded by this service.
   // This also clears the internal map of metrics as a biproduct, so if
   // comparing numbers of reports, the comparison should be from the last time
@@ -32,6 +38,13 @@ class TestRapporService : public RapporService {
   // This clears the internal map of metrics.
   void GetReports(RapporReports* reports);
 
+  // Gets the recorded sample/type for a |metric_name|, and returns whether the
+  // recorded metric was found. Limitation: if the metric was logged more than
+  // once, this will return the latest sample that was logged.
+  bool GetRecordedSampleForMetric(const std::string& metric_name,
+                                  std::string* sample,
+                                  RapporType* type);
+
   void set_is_incognito(bool is_incognito) { is_incognito_ = is_incognito; }
 
   TestingPrefServiceSimple* test_prefs() { return &test_prefs_; }
@@ -48,6 +61,14 @@ class TestRapporService : public RapporService {
   void ScheduleNextLogRotation(base::TimeDelta interval) override;
 
  private:
+  // Used to keep track of recorded RAPPOR samples.
+  struct RapporSample {
+    RapporType type;
+    std::string value;
+  };
+  typedef std::map<std::string, RapporSample> SamplesMap;
+  SamplesMap samples_;
+
   TestingPrefServiceSimple test_prefs_;
 
   // Holds a weak ref to the uploader_ object.