[Media Router] Assign each route a current cast mode if possible

chrome/browser/ui/webui/media_router/media_router_ui_unittest.cc

Index: chrome/browser/ui/webui/media_router/media_router_ui_unittest.cc
diff --git a/chrome/browser/ui/webui/media_router/media_router_ui_unittest.cc b/chrome/browser/ui/webui/media_router/media_router_ui_unittest.cc
index c70649dc5e8804224a968d4ff44599ca9cf329ab..8e21520ac61810855d79fa6e93221637204872df 100644
--- a/chrome/browser/ui/webui/media_router/media_router_ui_unittest.cc
+++ b/chrome/browser/ui/webui/media_router/media_router_ui_unittest.cc
@@ -8,6 +8,7 @@
 #include "base/memory/ptr_util.h"
 #include "base/strings/utf_string_conversions.h"
 #include "chrome/browser/media/router/media_route.h"
+#include "chrome/browser/media/router/media_source_helper.h"
 #include "chrome/browser/media/router/mock_media_router.h"
 #include "chrome/browser/media/router/route_request_result.h"
 #include "chrome/browser/media/router/test_helper.h"
@@ -217,10 +218,9 @@ TEST_F(MediaRouterUITest, UIMediaRoutesObserverFiltersNonDisplayRoutes) {
   routes.push_back(display_route_2);
 
   std::vector<MediaRoute> filtered_routes;
-  EXPECT_CALL(mock_callback, OnRoutesUpdated(_, _)).WillOnce(
-      SaveArg<0>(&filtered_routes));
-  observer->OnRoutesUpdated(routes,
-                            std::vector<MediaRoute::Id>());
+  EXPECT_CALL(mock_callback, OnRoutesUpdated(_, _))
+      .WillOnce(SaveArg<0>(&filtered_routes));
+  observer->OnRoutesUpdated(routes, std::vector<MediaRoute::Id>());
 
   ASSERT_EQ(2u, filtered_routes.size());
   EXPECT_TRUE(display_route_1.Equals(filtered_routes[0]));
@@ -261,10 +261,9 @@ TEST_F(MediaRouterUITest,
 
   std::vector<MediaRoute::Id> filtered_joinable_route_ids;
   // Save the filtered joinable routes.
-  EXPECT_CALL(mock_callback, OnRoutesUpdated(_, _)).WillOnce(
-      SaveArg<1>(&filtered_joinable_route_ids));
-  observer->OnRoutesUpdated(routes,
-                            joinable_route_ids);
+  EXPECT_CALL(mock_callback, OnRoutesUpdated(_, _))
+      .WillOnce(SaveArg<1>(&filtered_joinable_route_ids));
+  observer->OnRoutesUpdated(routes, joinable_route_ids);
 
   ASSERT_EQ(2u, filtered_joinable_route_ids.size());
   EXPECT_EQ(display_route_1.media_route_id(), filtered_joinable_route_ids[0]);
@@ -274,6 +273,105 @@ TEST_F(MediaRouterUITest,
   observer.reset();
 }
 
+TEST_F(MediaRouterUITest, UIMediaRoutesObserverAssignsCurrentCastModes) {
+  CreateMediaRouterUI(&profile_);
+  SessionID::id_type tab_id = SessionTabHelper::IdForTab(initiator_.get());
+  MediaSource media_source_1(MediaSourceForTab(tab_id));
+  MediaSource media_source_2("mediaSource");
+  MediaSource media_source_3(MediaSourceForDesktop());
+  MockRoutesUpdatedCallback mock_callback;
+  std::unique_ptr<MediaRouterUI::UIMediaRoutesObserver> observer(
+      new MediaRouterUI::UIMediaRoutesObserver(
+          &mock_router_, MediaSource::Id(),
+          base::Bind(&MediaRouterUI::OnRoutesUpdated,
+                     base::Unretained(media_router_ui_.get()))));
+
+  MediaRoute display_route_1("routeId1", media_source_1, "sinkId1", "desc 1",
+                             true, "", true);
+  MediaRoute non_display_route_1("routeId2", media_source_2, "sinkId2",
+                                 "desc 2", true, "", false);
+  MediaRoute display_route_2("routeId3", media_source_3, "sinkId2", "desc 2",
+                             true, "", true);
+  std::vector<MediaRoute> routes;
+  routes.push_back(display_route_1);
+  routes.push_back(non_display_route_1);
+  routes.push_back(display_route_2);
+
+  observer->OnRoutesUpdated(routes, std::vector<MediaRoute::Id>());
+
+  const auto& filtered_routes = media_router_ui_->routes();
+  ASSERT_EQ(2u, filtered_routes.size());
+  EXPECT_TRUE(display_route_1.Equals(filtered_routes[0]));
+  EXPECT_TRUE(filtered_routes[0].for_display());
+  EXPECT_TRUE(display_route_2.Equals(filtered_routes[1]));
+  EXPECT_TRUE(filtered_routes[1].for_display());
+
+  const auto& current_cast_modes = media_router_ui_->current_cast_modes();
+  ASSERT_EQ(2u, current_cast_modes.size());
+  auto cast_mode_entry =
+      current_cast_modes.find(display_route_1.media_route_id());
+  EXPECT_NE(end(current_cast_modes), cast_mode_entry);
+  EXPECT_EQ(MediaCastMode::TAB_MIRROR, cast_mode_entry->second);
+  cast_mode_entry =
+      current_cast_modes.find(non_display_route_1.media_route_id());
+  EXPECT_EQ(end(current_cast_modes), cast_mode_entry);
+  cast_mode_entry = current_cast_modes.find(display_route_2.media_route_id());
+  EXPECT_NE(end(current_cast_modes), cast_mode_entry);
+  EXPECT_EQ(MediaCastMode::DESKTOP_MIRROR, cast_mode_entry->second);
+
+  EXPECT_CALL(mock_router_, UnregisterMediaRoutesObserver(_)).Times(1);
+  observer.reset();
+}
+
+TEST_F(MediaRouterUITest, UIMediaRoutesObserverSkipsUnavailableCastModes) {
+  CreateMediaRouterUI(&profile_);
+  MediaSource media_source_1("mediaSource1");
+  MediaSource media_source_2("mediaSource2");
+  MediaSource media_source_3(MediaSourceForDesktop());
+  MockRoutesUpdatedCallback mock_callback;
+  std::unique_ptr<MediaRouterUI::UIMediaRoutesObserver> observer(
+      new MediaRouterUI::UIMediaRoutesObserver(
+          &mock_router_, MediaSource::Id(),
+          base::Bind(&MediaRouterUI::OnRoutesUpdated,
+                     base::Unretained(media_router_ui_.get()))));
+
+  MediaRoute display_route_1("routeId1", media_source_1, "sinkId1", "desc 1",
+                             true, "", true);
+  MediaRoute non_display_route_1("routeId2", media_source_2, "sinkId2",
+                                 "desc 2", true, "", false);
+  MediaRoute display_route_2("routeId3", media_source_3, "sinkId2", "desc 2",
+                             true, "", true);
+  std::vector<MediaRoute> routes;
+  routes.push_back(display_route_1);
+  routes.push_back(non_display_route_1);
+  routes.push_back(display_route_2);
+
+  observer->OnRoutesUpdated(routes, std::vector<MediaRoute::Id>());
+
+  const auto& filtered_routes = media_router_ui_->routes();
+  ASSERT_EQ(2u, filtered_routes.size());
+  EXPECT_TRUE(display_route_1.Equals(filtered_routes[0]));
+  EXPECT_TRUE(filtered_routes[0].for_display());
+  EXPECT_TRUE(display_route_2.Equals(filtered_routes[1]));
+  EXPECT_TRUE(filtered_routes[1].for_display());
+
+  const auto& current_cast_modes = media_router_ui_->current_cast_modes();
+  ASSERT_EQ(1u, current_cast_modes.size());
+  auto cast_mode_entry =
+      current_cast_modes.find(display_route_1.media_route_id());
+  // No observer for source "mediaSource1" means no cast mode for this route.
+  EXPECT_EQ(end(current_cast_modes), cast_mode_entry);
+  cast_mode_entry =
+      current_cast_modes.find(non_display_route_1.media_route_id());
+  EXPECT_EQ(end(current_cast_modes), cast_mode_entry);
+  cast_mode_entry = current_cast_modes.find(display_route_2.media_route_id());
+  EXPECT_NE(end(current_cast_modes), cast_mode_entry);
+  EXPECT_EQ(MediaCastMode::DESKTOP_MIRROR, cast_mode_entry->second);
+
+  EXPECT_CALL(mock_router_, UnregisterMediaRoutesObserver(_)).Times(1);
+  observer.reset();
+}
+
 TEST_F(MediaRouterUITest, GetExtensionNameExtensionPresent) {
   std::string id = "extensionid";
   GURL url = GURL("chrome-extension://" + id);