1 files changed, 50 insertions, 0 deletions
diff --git a/extractor/phrase_location_sampler_test.cc b/extractor/phrase_location_sampler_test.cc
new file mode 100644
index 00000000..e7520ce7
--- /dev/null
+++ b/extractor/phrase_location_sampler_test.cc
@@ -0,0 +1,50 @@
+#include <gtest/gtest.h>
+
+#include <memory>
+
+#include "mocks/mock_matchings_sampler.h"
+#include "mocks/mock_suffix_array_sampler.h"
+#include "phrase_location.h"
+#include "phrase_location_sampler.h"
+
+using namespace std;
+using namespace ::testing;
+
+namespace extractor {
+namespace {
+
+class MatchingsSamplerTest : public Test {
+ protected:
+  virtual void SetUp() {
+    matchings_sampler = make_shared<MockMatchingsSampler>();
+    suffix_array_sampler = make_shared<MockSuffixArraySampler>();
+
+    sampler = make_shared<PhraseLocationSampler>(
+        matchings_sampler, suffix_array_sampler);
+  }
+
+  shared_ptr<MockMatchingsSampler> matchings_sampler;
+  shared_ptr<MockSuffixArraySampler> suffix_array_sampler;
+  shared_ptr<PhraseLocationSampler> sampler;
+};
+
+TEST_F(MatchingsSamplerTest, TestSuffixArrayRange) {
+  vector<int> locations = {0, 1, 2, 3};
+  PhraseLocation location(0, 3), result(locations, 2);
+  unordered_set<int> blacklisted_sentence_ids;
+  EXPECT_CALL(*suffix_array_sampler, Sample(location, blacklisted_sentence_ids))
+      .WillOnce(Return(result));
+  EXPECT_EQ(result, sampler->Sample(location, blacklisted_sentence_ids));
+}
+
+TEST_F(MatchingsSamplerTest, TestMatchings) {
+  vector<int> locations = {0, 1, 2, 3};
+  PhraseLocation location(locations, 2), result(locations, 2);
+  unordered_set<int> blacklisted_sentence_ids;
+  EXPECT_CALL(*matchings_sampler, Sample(location, blacklisted_sentence_ids))
+      .WillOnce(Return(result));
+  EXPECT_EQ(result, sampler->Sample(location, blacklisted_sentence_ids));
+}
+
+}
+} // namespace extractor