Index: test/ReferrerMapping.cpp
diff --git a/test/ReferrerMapping.cpp b/test/ReferrerMapping.cpp
index 63d8e16511aec5b3bf7c2332c6a193a3034d9852..f645aedd351dd2c7fdc100d364d7920cc6c2facd 100644
--- a/test/ReferrerMapping.cpp
+++ b/test/ReferrerMapping.cpp
@@ -18,21 +18,55 @@
 #include <AdblockPlus.h>
 #include <gtest/gtest.h>
 
+typedef AdblockPlus::FilterEngine::ContentType ContentType;
+typedef AdblockPlus::ReferrerMapping::FrameIndicator FrameIndicator;
+
 TEST(ReferrerMappingTest, EmptyReferrerChain)
 {
   AdblockPlus::ReferrerMapping referrerMapping;
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("first");
+    referrerMapping.BuildFrameStructure("first");
   ASSERT_EQ(1, referrerChain.size());
-  ASSERT_EQ("first", referrerChain[0]);
+  EXPECT_EQ("first", referrerChain[0]);
+}
+
+TEST(ReferrerMappingTest, EmptyUrl)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  std::vector<std::string> referrerChain =
+    referrerMapping.BuildFrameStructure("");
+  EXPECT_EQ(0, referrerChain.size());
+
+  referrerMapping.Add("", "referrer", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerChain = referrerMapping.BuildFrameStructure("");
+  EXPECT_EQ(0, referrerChain.size());
+}
+
+TEST(ReferrerMappingTest, EmptyReferrer)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  referrerMapping.Add("first", "", FrameIndicator::FRAME_INDICATOR_FRAME);
+  std::vector<std::string> referrerChain =
+    referrerMapping.BuildFrameStructure("first");
+  ASSERT_EQ(1, referrerChain.size());
+  EXPECT_EQ("first", referrerChain[0]);
+}
+
+TEST(ReferrerMappingTest, NotAFrame)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  referrerMapping.Add("first", "", FrameIndicator::FRAME_INDICATOR_NOT_FRAME);
+  std::vector<std::string> referrerChain =
+    referrerMapping.BuildFrameStructure("first");
+  ASSERT_EQ(0, referrerChain.size());
 }
 
 TEST(ReferrerMappingTest, TwoElementReferrerChain)
 {
   AdblockPlus::ReferrerMapping referrerMapping;
-  referrerMapping.Add("second", "first");
+  referrerMapping.Add("second", "first", FrameIndicator::FRAME_INDICATOR_FRAME);
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("second");
+    referrerMapping.BuildFrameStructure("second");
   ASSERT_EQ(2, referrerChain.size());
   ASSERT_EQ("first", referrerChain[0]);
   ASSERT_EQ("second", referrerChain[1]);
@@ -41,17 +75,17 @@ TEST(ReferrerMappingTest, TwoElementReferrerChain)
 TEST(ReferrerMappingTest, TenElementReferrerChain)
 {
   AdblockPlus::ReferrerMapping referrerMapping;
-  referrerMapping.Add("second", "first");
-  referrerMapping.Add("third", "second");
-  referrerMapping.Add("fourth", "third");
-  referrerMapping.Add("fifth", "fourth");
-  referrerMapping.Add("sixth", "fifth");
-  referrerMapping.Add("seventh", "sixth");
-  referrerMapping.Add("eighth", "seventh");
-  referrerMapping.Add("ninth", "eighth");
-  referrerMapping.Add("tenth", "ninth");
+  referrerMapping.Add("second", "first", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("third", "second", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fourth", "third", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fifth", "fourth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("sixth", "fifth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("seventh", "sixth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("eighth", "seventh", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("ninth", "eighth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("tenth", "ninth", FrameIndicator::FRAME_INDICATOR_FRAME);
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("tenth");
+    referrerMapping.BuildFrameStructure("tenth");
   ASSERT_EQ(10, referrerChain.size());
   ASSERT_EQ("first", referrerChain[0]);
   ASSERT_EQ("second", referrerChain[1]);
@@ -68,14 +102,15 @@ TEST(ReferrerMappingTest, TenElementReferrerChain)
 TEST(ReferrerMappingTest, CacheOnlyFiveUrls)
 {
   AdblockPlus::ReferrerMapping referrerMapping(5);
-  referrerMapping.Add("second", "first");
-  referrerMapping.Add("third", "second");
-  referrerMapping.Add("fourth", "third");
-  referrerMapping.Add("fifth", "fourth");
-  referrerMapping.Add("sixth", "fifth");
-  referrerMapping.Add("seventh", "sixth");
+  std::string outOfCacheUrl = "first";
+  referrerMapping.Add("second", outOfCacheUrl, FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("third", "second", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fourth", "third", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("fifth", "fourth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("sixth", "fifth", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("seventh", "sixth", FrameIndicator::FRAME_INDICATOR_FRAME);
   std::vector<std::string> referrerChain =
-    referrerMapping.BuildReferrerChain("seventh");
+    referrerMapping.BuildFrameStructure("seventh");
   ASSERT_EQ(6, referrerChain.size());
   ASSERT_EQ("second", referrerChain[0]);
   ASSERT_EQ("third", referrerChain[1]);
@@ -84,3 +119,26 @@ TEST(ReferrerMappingTest, CacheOnlyFiveUrls)
   ASSERT_EQ("sixth", referrerChain[4]);
   ASSERT_EQ("seventh", referrerChain[5]);
 }
+
+TEST(ReferrerMappingTest, OnlyFrames)
+{
+  AdblockPlus::ReferrerMapping referrerMapping;
+  referrerMapping.Add("iframe1", "root-page", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("not a frame", "iframe1", FrameIndicator::FRAME_INDICATOR_NOT_FRAME);
+  referrerMapping.Add("iframe2", "not a frame", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerMapping.Add("media", "iframe2", FrameIndicator::FRAME_INDICATOR_NOT_FRAME);
+
+  std::vector<std::string> referrerChain = referrerMapping.BuildFrameStructure("media");
+  ASSERT_EQ(3, referrerChain.size());
+  EXPECT_EQ("root-page", referrerChain[0]);
+  EXPECT_EQ("iframe1", referrerChain[1]);
+  EXPECT_EQ("iframe2", referrerChain[2]);
+
+  referrerMapping.Add("iframe3", "media", FrameIndicator::FRAME_INDICATOR_FRAME);
+  referrerChain = referrerMapping.BuildFrameStructure("iframe3");
+  ASSERT_EQ(4, referrerChain.size());
+  EXPECT_EQ("root-page", referrerChain[0]);
+  EXPECT_EQ("iframe1", referrerChain[1]);
+  EXPECT_EQ("iframe2", referrerChain[2]);
+  EXPECT_EQ("iframe3", referrerChain[3]);
+}
\ No newline at end of file