[Media Router] Add a presentation id to MediaRoute mapping in the MR

chrome/browser/media/router/media_router_base_unittest.cc

Index: chrome/browser/media/router/media_router_base_unittest.cc
diff --git a/chrome/browser/media/router/media_router_base_unittest.cc b/chrome/browser/media/router/media_router_base_unittest.cc
index 4b942cfcd1f2134bbe52d1fcac6e60c930fb5b54..1754ce9bd6a7443fab30e10329b1b7da63b0645b 100644
--- a/chrome/browser/media/router/media_router_base_unittest.cc
+++ b/chrome/browser/media/router/media_router_base_unittest.cc
@@ -39,6 +39,16 @@ class MockMediaRouterBase : public MockMediaRouter {
   std::vector<MediaRoute> GetCurrentRoutes() const override {
     return MediaRouterBase::GetCurrentRoutes();
   }
+
+  const MediaRoute* GetPresentationRoute(
+      const std::string& presentation_id) const override {
+    return MediaRouterBase::GetPresentationRoute(presentation_id);
+  }
+
+  void SetPresentationId(const MediaRoute::Id& route_id,
+                         const std::string& presentation_id) override {
+    MediaRouterBase::SetPresentationId(route_id, presentation_id);
+  }
 };
 
 class MediaRouterBaseTest : public testing::Test {
@@ -136,4 +146,98 @@ TEST_F(MediaRouterBaseTest, GetCurrentRoutes) {
   EXPECT_TRUE(router_.GetCurrentRoutes().empty());
 }
 
+TEST_F(MediaRouterBaseTest, GetPresentationRoute) {
+  MediaSource source1("source_1");
+  MediaSource source2("source_2");
+  MediaRoute route1("route_1", source1, "sink_1", "", false, "", false);
+  MediaRoute route2("route_2", source2, "sink_2", "", true, "", false);
+  std::string presentation_id1("presentation_id1");
+  std::string presentation_id2("presentation_id2");
+  std::string presentation_id3("presentation_id3");
+  route1.set_presentation_id(presentation_id1);
+  route2.set_presentation_id(presentation_id2);
+  std::vector<MediaRoute> routes = {route1, route2};
+  std::vector<MediaRoute::Id> joinable_route_ids = {"route_1"};
+
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id1));
+
+  routes_observer_->OnRoutesUpdated(routes, joinable_route_ids);
+  auto* actual_route1 = router_.GetPresentationRoute(presentation_id1);
+  auto* actual_route2 = router_.GetPresentationRoute(presentation_id2);
+
+  EXPECT_TRUE(actual_route1 && actual_route1->Equals(route1));
+  EXPECT_TRUE(actual_route2 && actual_route2->Equals(route2));
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id3));
+
+  routes.erase(routes.begin());
+  routes_observer_->OnRoutesUpdated(routes, joinable_route_ids);
+
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id1));
+}
+
+TEST_F(MediaRouterBaseTest, OnRoutes) {
+  MediaSource source1("source_1");
+  MediaSource source2("source_2");
+  MediaRoute route1("route_1", source1, "sink_1", "", false, "", false);
+  MediaRoute route2("route_2", source2, "sink_2", "", true, "", false);
+  std::string presentation_id1("presentation_id1");
+  std::string presentation_id2("presentation_id2");
+  std::string presentation_id3("presentation_id3");
+  std::vector<MediaRoute> routes = {route1, route2};
+  std::vector<MediaRoute::Id> joinable_route_ids = {"route_1"};
+
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id1));
+
+  routes_observer_->OnRoutesUpdated(routes, joinable_route_ids);
+  router_.SetPresentationId(route1.media_route_id(), presentation_id1);
+  router_.SetPresentationId(route2.media_route_id(), presentation_id2);
+  router_.SetPresentationId("route3", presentation_id3);
+
+  auto* actual_route1 = router_.GetPresentationRoute(presentation_id1);
+  auto* actual_route2 = router_.GetPresentationRoute(presentation_id2);
+
+  EXPECT_TRUE(actual_route1 && actual_route1->Equals(route1));
+  EXPECT_TRUE(actual_route2 && actual_route2->Equals(route2));
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id3));
+
+  routes.erase(routes.begin());
+  routes_observer_->OnRoutesUpdated(routes, joinable_route_ids);
+
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id1));
+}
+
+TEST_F(MediaRouterBaseTest, OnRoutesUpdated) {
+  MediaSource source1("source_1");
+  MediaSource source2("source_2");
+  MediaRoute route1("route_1", source1, "sink_1", "", false, "", false);
+  MediaRoute route2("route_2", source2, "sink_2", "", true, "", false);
+  std::string presentation_id1("presentation_id1");
+  std::string presentation_id2("presentation_id2");
+  std::string presentation_id3("presentation_id3");
+  std::string presentation_id4("presentation_id4");
+
+  std::vector<MediaRoute> routes1 = {route1, route2};
+  std::vector<MediaRoute::Id> joinable_route_ids = {"route_1"};
+  routes_observer_->OnRoutesUpdated(routes1, joinable_route_ids);
+  router_.SetPresentationId(route2.media_route_id(), presentation_id2);
+
+  std::vector<MediaRoute> routes2 = {route1, route2};
+  routes_observer_->OnRoutesUpdated(routes2, joinable_route_ids);
+
+  // Preserve old presentation id.
+  auto* actual_route2 = router_.GetPresentationRoute(presentation_id2);
+  EXPECT_TRUE(actual_route2 && actual_route2->Equals(route2));
+
+  // Use new presentation id.
+  route1.set_presentation_id(presentation_id3);
+  route2.set_presentation_id(presentation_id4);
+
+  std::vector<MediaRoute> routes3 = {route1, route2};
+  routes_observer_->OnRoutesUpdated(routes3, joinable_route_ids);
+
+  EXPECT_FALSE(router_.GetPresentationRoute(presentation_id2));
+  auto* actual_route4 = router_.GetPresentationRoute(presentation_id4);
+  EXPECT_TRUE(actual_route4 && actual_route4->Equals(route2));
+}
+
 }  // namespace media_router