[Media Router] Support dual discovery

chrome/browser/media/router/discovery/mdns/cast_media_sink_service_unittest.cc

 // Copyright 2017 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "chrome/browser/media/router/discovery/mdns/cast_media_sink_service.h"
 #include "base/strings/stringprintf.h"
 #include "base/test/mock_callback.h"
 #include "base/timer/mock_timer.h"
 #include "chrome/browser/media/router/discovery/mdns/mock_dns_sd_registry.h"
 #include "chrome/browser/media/router/test_helper.h"
 #include "chrome/test/base/testing_profile.h"
 #include "components/cast_channel/cast_socket.h"
 #include "components/cast_channel/cast_socket_service.h"
 #include "components/cast_channel/cast_test_util.h"
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 using ::testing::_;
+using ::testing::Invoke;
 using ::testing::Return;
 using ::testing::SaveArg;
+using ::testing::WithArgs;
+
+using cast_channel::CastSocket;
+using cast_channel::ChannelError;
 
 namespace {
 
-net::IPEndPoint CreateIPEndPoint(int num) {
+constexpr char SERVICE_TYPE[] = "_googlecast._tcp.local";
+
+// Test parameter structure holding mock strings for media sink and mDNS
+// service.
+struct TestSinkParams {
+  std::string unique_id;
+  std::string model_name;
+  std::string friendly_name;
   net::IPAddress ip_address;
-  DCHECK(ip_address.AssignFromIPLiteral(
-      base::StringPrintf("192.168.0.10%d", num)));
-  return net::IPEndPoint(ip_address, 8009 + num);
-}
+  net::IPEndPoint ip_endpoint;
+  int channel_id;
+  bool audio_only;
 
-media_router::DnsSdService CreateDnsService(int num) {
-  net::IPEndPoint ip_endpoint = CreateIPEndPoint(num);
+  TestSinkParams() = default;
+
+  explicit TestSinkParams(int num) {
+    unique_id = base::StringPrintf("service %d", num);
+    friendly_name = base::StringPrintf("friendly name %d", num);
+    model_name = base::StringPrintf("model name %d", num);
+    DCHECK(ip_address.AssignFromIPLiteral(
+        base::StringPrintf("192.168.0.10%d", num)));
+    ip_endpoint = net::IPEndPoint(
+        ip_address, media_router::CastMediaSinkService::kCastControlPort + num);
+    channel_id = num;
+    audio_only = false;
+  }
+};
+
+media_router::DnsSdService CreateDnsService(const TestSinkParams& params) {
   media_router::DnsSdService service;
-  service.service_name =
-      "_myDevice." +
-      std::string(media_router::CastMediaSinkService::kCastServiceType);
-  service.ip_address = ip_endpoint.address().ToString();
-  service.service_host_port = ip_endpoint.ToString();
-  service.service_data.push_back(base::StringPrintf("id=service %d", num));
-  service.service_data.push_back(
-      base::StringPrintf("fn=friendly name %d", num));
-  service.service_data.push_back(base::StringPrintf("md=model name %d", num));
-
+  service.service_name = "_myDevice." + std::string(SERVICE_TYPE);
+  service.ip_address = params.ip_address.ToString();
+  service.service_host_port = params.ip_endpoint.ToString();
+  service.service_data.push_back("id=" + params.unique_id);
+  service.service_data.push_back("fn=" + params.friendly_name);
+  service.service_data.push_back("md=" + params.model_name);
   return service;
 }
 
+media_router::MediaSinkInternal CreateCastSink(const TestSinkParams& params) {
+  media_router::MediaSink sink(params.unique_id, params.friendly_name,
+                               media_router::MediaSink::IconType::CAST);
+
+  media_router::CastSinkExtraData extra_data;
+  extra_data.ip_address = params.ip_address;
+  extra_data.port = params.ip_endpoint.port();
+  extra_data.model_name = params.model_name;
+  extra_data.cast_channel_id = params.channel_id;
+  extra_data.capabilities = cast_channel::CastDeviceCapability::AUDIO_OUT |
+                            cast_channel::CastDeviceCapability::VIDEO_OUT;
+  return media_router::MediaSinkInternal(sink, extra_data);
+}
+
+media_router::MediaSinkInternal CreateDialSink(const TestSinkParams& params) {
+  media_router::MediaSink sink(params.unique_id, params.friendly_name,
+                               media_router::MediaSink::IconType::CAST);
+
+  media_router::DialSinkExtraData extra_data;
+  extra_data.ip_address = params.ip_address;
+  extra_data.model_name = params.model_name;
+  return media_router::MediaSinkInternal(sink, extra_data);
+}
+
 void VerifyMediaSinkInternal(const media_router::MediaSinkInternal& cast_sink,
-                             const media_router::DnsSdService& service,
-                             int channel_id,
-                             bool audio_only) {
-  std::string id = base::StringPrintf("service %d", channel_id);
-  std::string name = base::StringPrintf("friendly name %d", channel_id);
-  std::string model_name = base::StringPrintf("model name %d", channel_id);
-  EXPECT_EQ(id, cast_sink.sink().id());
-  EXPECT_EQ(name, cast_sink.sink().name());
-  EXPECT_EQ(model_name, cast_sink.cast_data().model_name);
-  EXPECT_EQ(service.ip_address, cast_sink.cast_data().ip_address.ToString());
+                             const TestSinkParams& params) {
+  EXPECT_EQ(params.unique_id, cast_sink.sink().id());
+  EXPECT_EQ(params.friendly_name, cast_sink.sink().name());
+  EXPECT_EQ(params.model_name, cast_sink.cast_data().model_name);
+  EXPECT_EQ(params.ip_address, cast_sink.cast_data().ip_address);
   EXPECT_TRUE(cast_channel::CastDeviceCapability::AUDIO_OUT &
               cast_sink.cast_data().capabilities);
-  EXPECT_EQ(!audio_only, cast_channel::CastDeviceCapability::VIDEO_OUT &
-                             cast_sink.cast_data().capabilities);
-  EXPECT_EQ(channel_id, cast_sink.cast_data().cast_channel_id);
+  EXPECT_EQ(!params.audio_only, cast_channel::CastDeviceCapability::VIDEO_OUT &
+                                    cast_sink.cast_data().capabilities);
+  EXPECT_EQ(params.channel_id, cast_sink.cast_data().cast_channel_id);
 }
 
 }  // namespace
 
 namespace media_router {
 
 class MockCastSocketService : public cast_channel::CastSocketService {
  public:
   MOCK_METHOD4(OpenSocket,
                int(const net::IPEndPoint& ip_endpoint,
                    net::NetLog* net_log,
                    const cast_channel::CastSocket::OnOpenCallback& open_cb,
                    cast_channel::CastSocket::Observer* observer));
   MOCK_CONST_METHOD1(GetSocket, cast_channel::CastSocket*(int channel_id));
 
  private:
   ~MockCastSocketService() {}
 };
 
 class CastMediaSinkServiceTest : public ::testing::Test {
  public:
   CastMediaSinkServiceTest()
       : mock_cast_socket_service_(new MockCastSocketService()),
         media_sink_service_(
             new CastMediaSinkService(mock_sink_discovered_cb_.Get(),
                                      mock_cast_socket_service_.get())),
-        test_dns_sd_registry_(media_sink_service_.get()) {}
+        test_dns_sd_registry_(media_sink_service_.get()) {
+    for (int i = 1; i <= 3; i++) {
+      params_[i] = TestSinkParams(i);
+      services_[i] = CreateDnsService(params_[i]);
+      cast_sinks_[i] = CreateCastSink(params_[i]);
+      dial_sinks_[i] = CreateDialSink(params_[i]);
+    }
+  }
 
   void SetUp() override {
     auto mock_timer = base::MakeUnique<base::MockTimer>(
         true /*retain_user_task*/, false /*is_repeating*/);
     mock_timer_ = mock_timer.get();
     media_sink_service_->SetTimerForTest(std::move(mock_timer));
   }
 
  protected:
   const content::TestBrowserThreadBundle thread_bundle_;
   TestingProfile profile_;
 
   base::MockCallback<MediaSinkService::OnSinksDiscoveredCallback>
       mock_sink_discovered_cb_;
   scoped_refptr<MockCastSocketService> mock_cast_socket_service_;
   scoped_refptr<CastMediaSinkService> media_sink_service_;
   MockDnsSdRegistry test_dns_sd_registry_;
   base::MockTimer* mock_timer_;
 
+  TestSinkParams params_[4];
+  DnsSdService services_[4];
+  MediaSinkInternal cast_sinks_[4];
+  MediaSinkInternal dial_sinks_[4];
+  cast_channel::MockCastSocket mock_sockets_[4];
+
   DISALLOW_COPY_AND_ASSIGN(CastMediaSinkServiceTest);
 };
 
 TEST_F(CastMediaSinkServiceTest, TestReStartAfterStop) {
   EXPECT_CALL(test_dns_sd_registry_, AddObserver(media_sink_service_.get()))
       .Times(2);
   EXPECT_CALL(test_dns_sd_registry_, RegisterDnsSdListener(_)).Times(2);
   EXPECT_FALSE(mock_timer_->IsRunning());
   media_sink_service_->SetDnsSdRegistryForTest(&test_dns_sd_registry_);
   media_sink_service_->Start();
   EXPECT_TRUE(mock_timer_->IsRunning());
 
   EXPECT_CALL(test_dns_sd_registry_, RemoveObserver(media_sink_service_.get()));
   EXPECT_CALL(test_dns_sd_registry_, UnregisterDnsSdListener(_));
   media_sink_service_->Stop();
 
-  mock_timer_ =
-      new base::MockTimer(true /*retain_user_task*/, false /*is_repeating*/);
-  media_sink_service_->SetTimerForTest(base::WrapUnique(mock_timer_));
+  auto mock_timer = base::MakeUnique<base::MockTimer>(true /*retain_user_task*/,
+                                                      false /*is_repeating*/);
+  mock_timer_ = mock_timer.get();
+  media_sink_service_->SetTimerForTest(std::move(mock_timer));
   media_sink_service_->SetDnsSdRegistryForTest(&test_dns_sd_registry_);
   media_sink_service_->Start();
   EXPECT_TRUE(mock_timer_->IsRunning());
 }
 
 TEST_F(CastMediaSinkServiceTest, TestMultipleStartAndStop) {
   EXPECT_CALL(test_dns_sd_registry_, AddObserver(media_sink_service_.get()));
   EXPECT_CALL(test_dns_sd_registry_, RegisterDnsSdListener(_));
   media_sink_service_->SetDnsSdRegistryForTest(&test_dns_sd_registry_);
   media_sink_service_->Start();
   media_sink_service_->Start();
   EXPECT_TRUE(mock_timer_->IsRunning());
 
   EXPECT_CALL(test_dns_sd_registry_, RemoveObserver(media_sink_service_.get()));
   EXPECT_CALL(test_dns_sd_registry_, UnregisterDnsSdListener(_));
   media_sink_service_->Stop();
   media_sink_service_->Stop();
 }
 
 TEST_F(CastMediaSinkServiceTest, TestOnChannelOpenedOnIOThread) {
-  DnsSdService service = CreateDnsService(1);
-  cast_channel::MockCastSocket socket;
   EXPECT_CALL(*mock_cast_socket_service_, GetSocket(1))
-      .WillOnce(Return(&socket));
+      .WillOnce(Return(&mock_sockets_[1]));
 
-  media_sink_service_->current_services_.push_back(service);
+  media_sink_service_->current_service_ip_endpoints_.insert(
+      params_[1].ip_endpoint);
   media_sink_service_->OnChannelOpenedOnIOThread(
-      service, 1, cast_channel::ChannelError::NONE);
+      cast_sinks_[1], 1, cast_channel::ChannelError::NONE);
   // Invoke CastMediaSinkService::OnChannelOpenedOnUIThread on the UI thread.
   base::RunLoop().RunUntilIdle();
 
-  // Verify sink content
-  EXPECT_EQ(size_t(1), media_sink_service_->current_sinks_.size());
-  for (const auto& sink_it : media_sink_service_->current_sinks_)
-    VerifyMediaSinkInternal(sink_it, service, 1, false);
+  // Verify sink content.
+  EXPECT_EQ(size_t(1), media_sink_service_->current_sinks_by_mdns_map_.size());
+  for (const auto& sink_it : media_sink_service_->current_sinks_by_mdns_map_)
+    EXPECT_EQ(cast_sinks_[1], sink_it.second);
 }
 
 TEST_F(CastMediaSinkServiceTest, TestMultipleOnChannelOpenedOnIOThread) {
-  DnsSdService service1 = CreateDnsService(1);
-  DnsSdService service2 = CreateDnsService(2);
-  DnsSdService service3 = CreateDnsService(3);
-
-  cast_channel::MockCastSocket socket2;
-  cast_channel::MockCastSocket socket3;
   // Fail to open channel 1.
   EXPECT_CALL(*mock_cast_socket_service_, GetSocket(1))
       .WillOnce(Return(nullptr));
   EXPECT_CALL(*mock_cast_socket_service_, GetSocket(2))
-      .WillOnce(Return(&socket2));
+      .WillOnce(Return(&mock_sockets_[2]));
   EXPECT_CALL(*mock_cast_socket_service_, GetSocket(3))
-      .WillOnce(Return(&socket2));
+      .WillOnce(Return(&mock_sockets_[3]));
 
   // Current round of Dns discovery finds service1 and service 2.
-  media_sink_service_->current_services_.push_back(service1);
-  media_sink_service_->current_services_.push_back(service2);
+  media_sink_service_->current_service_ip_endpoints_.insert(
+      params_[1].ip_endpoint);
+  media_sink_service_->current_service_ip_endpoints_.insert(
+      params_[2].ip_endpoint);
   media_sink_service_->OnChannelOpenedOnIOThread(
-      service1, 1, cast_channel::ChannelError::NONE);
+      cast_sinks_[1], 1, cast_channel::ChannelError::NONE);
   media_sink_service_->OnChannelOpenedOnIOThread(
-      service2, 2, cast_channel::ChannelError::NONE);
+      cast_sinks_[2], 2, cast_channel::ChannelError::NONE);
   media_sink_service_->OnChannelOpenedOnIOThread(
-      service3, 3, cast_channel::ChannelError::NONE);
+      cast_sinks_[3], 3, cast_channel::ChannelError::NONE);
   // Invoke CastMediaSinkService::OnChannelOpenedOnUIThread on the UI thread.
   base::RunLoop().RunUntilIdle();
 
-  // Verify sink content
-  EXPECT_EQ(size_t(1), media_sink_service_->current_sinks_.size());
-  for (const auto& sink_it : media_sink_service_->current_sinks_)
-    VerifyMediaSinkInternal(sink_it, service2, 2, false);
+  // Verify sink content.
+  EXPECT_EQ(size_t(1), media_sink_service_->current_sinks_by_mdns_map_.size());
+  for (const auto& sink_it : media_sink_service_->current_sinks_by_mdns_map_)
+    EXPECT_EQ(cast_sinks_[2], sink_it.second);
 }
 
 TEST_F(CastMediaSinkServiceTest, TestOnDnsSdEvent) {
-  DnsSdService service1 = CreateDnsService(1);
-  DnsSdService service2 = CreateDnsService(2);
-  net::IPEndPoint ip_endpoint1 = CreateIPEndPoint(1);
-  net::IPEndPoint ip_endpoint2 = CreateIPEndPoint(2);
+  // Add dns services.
+  DnsSdRegistry::DnsSdServiceList service_list{services_[1], services_[2]};
 
-  // Add dns services.
-  DnsSdRegistry::DnsSdServiceList service_list{service1, service2};
+  // Channel 1, 2 opened.
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[1].ip_endpoint, _, _, _))
+      .WillOnce(DoAll(
+          WithArgs<2>(Invoke([&](const CastSocket::OnOpenCallback& callback) {
+            callback.Run(1, ChannelError::NONE);
+          })),
+          Return(1)));
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[2].ip_endpoint, _, _, _))
+      .WillOnce(DoAll(
+          WithArgs<2>(Invoke([&](const CastSocket::OnOpenCallback& callback) {
+            callback.Run(2, ChannelError::NONE);
+          })),
+          Return(2)));
 
-  cast_channel::CastSocket::OnOpenCallback callback1;
-  cast_channel::CastSocket::OnOpenCallback callback2;
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint1, _, _, _))
-      .WillOnce(DoAll(SaveArg<2>(&callback1), Return(1)));
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint2, _, _, _))
-      .WillOnce(DoAll(SaveArg<2>(&callback2), Return(2)));
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(1))
+      .WillOnce(Return(&mock_sockets_[1]));
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(2))
+      .WillOnce(Return(&mock_sockets_[2]));
 
   // Invoke CastSocketService::OpenSocket on the IO thread.
-  media_sink_service_->OnDnsSdEvent(CastMediaSinkService::kCastServiceType,
-                                    service_list);
+  media_sink_service_->OnDnsSdEvent(SERVICE_TYPE, service_list);
   base::RunLoop().RunUntilIdle();
 
-  cast_channel::MockCastSocket socket1;
-  cast_channel::MockCastSocket socket2;
-
-  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(1))
-      .WillOnce(Return(&socket1));
-  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(2))
-      .WillOnce(Return(&socket2));
-
-  callback1.Run(1, cast_channel::ChannelError::NONE);
-  callback2.Run(2, cast_channel::ChannelError::NONE);
-
   // Invoke CastMediaSinkService::OnChannelOpenedOnUIThread on the UI thread.
   base::RunLoop().RunUntilIdle();
-  // Verify sink content
-  EXPECT_EQ(size_t(2), media_sink_service_->current_sinks_.size());
+  // Verify sink content.
+  EXPECT_EQ(size_t(2), media_sink_service_->current_sinks_by_mdns_map_.size());
 }
 
 TEST_F(CastMediaSinkServiceTest, TestMultipleOnDnsSdEvent) {
-  DnsSdService service1 = CreateDnsService(1);
-  DnsSdService service2 = CreateDnsService(2);
-  DnsSdService service3 = CreateDnsService(3);
-  net::IPEndPoint ip_endpoint1 = CreateIPEndPoint(1);
-  net::IPEndPoint ip_endpoint2 = CreateIPEndPoint(2);
-  net::IPEndPoint ip_endpoint3 = CreateIPEndPoint(3);
+  // 1st round finds service 1 & 2.
+  DnsSdRegistry::DnsSdServiceList service_list1{services_[1], services_[2]};
+  media_sink_service_->OnDnsSdEvent(SERVICE_TYPE, service_list1);
 
-  // 1st round finds service 1 & 2.
-  DnsSdRegistry::DnsSdServiceList service_list1{service1, service2};
-  media_sink_service_->OnDnsSdEvent(CastMediaSinkService::kCastServiceType,
-                                    service_list1);
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(1))
+      .WillOnce(Return(&mock_sockets_[1]));
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(2))
+      .WillOnce(Return(&mock_sockets_[2]));
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(3))
+      .WillOnce(Return(&mock_sockets_[3]));
 
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint1, _, _, _));
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint2, _, _, _));
+  // Channel 2 opened.
+  CastSocket::OnOpenCallback callback1;
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[1].ip_endpoint, _, _, _))
+      .WillOnce(DoAll(SaveArg<2>(&callback1), Return(1)));
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[2].ip_endpoint, _, _, _))
+      .WillOnce(DoAll(
+          WithArgs<2>(Invoke([&](const CastSocket::OnOpenCallback& callback) {
+            callback.Run(2, ChannelError::NONE);
+          })),
+          Return(2)));
   base::RunLoop().RunUntilIdle();
 
-  // Channel 2 opened.
-  media_sink_service_->OnChannelOpenedOnUIThread(service2, 2, false);
+  // 2nd round finds service 2 & 3.
+  DnsSdRegistry::DnsSdServiceList service_list2{services_[2], services_[3]};
+  media_sink_service_->OnDnsSdEvent(SERVICE_TYPE, service_list2);
 
-  // 2nd round finds service 2 & 3.
-  DnsSdRegistry::DnsSdServiceList service_list2{service2, service3};
-  media_sink_service_->OnDnsSdEvent(CastMediaSinkService::kCastServiceType,
-                                    service_list2);
-
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint2, _, _, _));
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint3, _, _, _));
+  // Channel 3 opened.
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[2].ip_endpoint, _, _, _));
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[3].ip_endpoint, _, _, _))
+      .WillOnce(DoAll(
+          WithArgs<2>(Invoke([&](const CastSocket::OnOpenCallback& callback) {
+            callback.Run(3, ChannelError::NONE);
+          })),
+          Return(3)));
   base::RunLoop().RunUntilIdle();
 
-  // Channel 1 and 3 opened.
-  media_sink_service_->OnChannelOpenedOnUIThread(service1, 1, false);
-  media_sink_service_->OnChannelOpenedOnUIThread(service3, 3, false);
+  // Channel 1 opened.
+  callback1.Run(1, ChannelError::NONE);
 
-  EXPECT_EQ(size_t(1), media_sink_service_->current_sinks_.size());
-  for (const auto& sink_it : media_sink_service_->current_sinks_)
-    VerifyMediaSinkInternal(sink_it, service3, 3, false);
+  // Verify sink content.
+  EXPECT_EQ(size_t(1), media_sink_service_->current_sinks_by_mdns_map_.size());
+  for (const auto& sink_it : media_sink_service_->current_sinks_by_mdns_map_)
+    VerifyMediaSinkInternal(sink_it.second, params_[3]);
 }
 
 TEST_F(CastMediaSinkServiceTest, TestTimer) {
-  DnsSdService service1 = CreateDnsService(1);
-  DnsSdService service2 = CreateDnsService(2);
-  net::IPEndPoint ip_endpoint1 = CreateIPEndPoint(1);
-  net::IPEndPoint ip_endpoint2 = CreateIPEndPoint(2);
-
   EXPECT_FALSE(mock_timer_->IsRunning());
   // finds service 1 & 2.
-  DnsSdRegistry::DnsSdServiceList service_list1{service1, service2};
-  media_sink_service_->OnDnsSdEvent(CastMediaSinkService::kCastServiceType,
-                                    service_list1);
+  DnsSdRegistry::DnsSdServiceList service_list1{services_[1], services_[2]};
+  media_sink_service_->OnDnsSdEvent(SERVICE_TYPE, service_list1);
 
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint1, _, _, _));
-  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint2, _, _, _));
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[1].ip_endpoint, _, _, _));
+  EXPECT_CALL(*mock_cast_socket_service_,
+              OpenSocket(params_[2].ip_endpoint, _, _, _));
   base::RunLoop().RunUntilIdle();
 
   // Channel 2 is opened.
-  media_sink_service_->OnChannelOpenedOnUIThread(service2, 2, false);
+  media_sink_service_->OnChannelOpenedOnUIThread(cast_sinks_[2]);
 
   std::vector<MediaSinkInternal> sinks;
   EXPECT_CALL(mock_sink_discovered_cb_, Run(_)).WillOnce(SaveArg<0>(&sinks));
 
   // Fire timer.
   mock_timer_->Fire();
   EXPECT_EQ(size_t(1), sinks.size());
-  VerifyMediaSinkInternal(sinks[0], service2, 2, false);
+  VerifyMediaSinkInternal(sinks[0], params_[2]);
 
   EXPECT_FALSE(mock_timer_->IsRunning());
   // Channel 1 is opened and timer is restarted.
-  media_sink_service_->OnChannelOpenedOnUIThread(service1, 1, false);
+  media_sink_service_->OnChannelOpenedOnUIThread(cast_sinks_[1]);
   EXPECT_TRUE(mock_timer_->IsRunning());
 }
 
+TEST_F(CastMediaSinkServiceTest, TestOnDialSinkAdded) {
+  net::IPEndPoint ip_endpoint1(params_[1].ip_address,
+                               CastMediaSinkService::kCastControlPort);
+  net::IPEndPoint ip_endpoint2(params_[2].ip_address,
+                               CastMediaSinkService::kCastControlPort);
+
+  // Channel 1, 2 opened.
+  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint1, _, _, _))
+      .WillOnce(DoAll(
+          WithArgs<2>(Invoke([&](const CastSocket::OnOpenCallback& callback) {
+            callback.Run(1, ChannelError::NONE);
+          })),
+          Return(1)));
+  EXPECT_CALL(*mock_cast_socket_service_, OpenSocket(ip_endpoint2, _, _, _))
+      .WillOnce(DoAll(
+          WithArgs<2>(Invoke([&](const CastSocket::OnOpenCallback& callback) {
+            callback.Run(2, ChannelError::NONE);
+          })),
+          Return(2)));
+
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(1))
+      .WillOnce(Return(&mock_sockets_[1]));
+  EXPECT_CALL(*mock_cast_socket_service_, GetSocket(2))
+      .WillOnce(Return(&mock_sockets_[2]));
+
+  // Invoke CastSocketService::OpenSocket on the IO thread.
+  media_sink_service_->OnMediaSinkAdded(dial_sinks_[1]);
+  base::RunLoop().RunUntilIdle();
+
+  // Invoke CastSocketService::OpenSocket on the IO thread.
+  media_sink_service_->OnMediaSinkAdded(dial_sinks_[2]);
+  base::RunLoop().RunUntilIdle();
+  // Verify sink content.
+  EXPECT_EQ(size_t(2), media_sink_service_->current_sinks_by_dial_map_.size());
+}
+
+TEST_F(CastMediaSinkServiceTest, TestOnFetchCompleted) {
+  std::vector<MediaSinkInternal> sinks;
+  EXPECT_CALL(mock_sink_discovered_cb_, Run(_)).WillOnce(SaveArg<0>(&sinks));
+
+  // Cast sink 1, 2 from mDNS discovery
+  media_sink_service_->current_sinks_by_mdns_map_[params_[1].ip_endpoint] =
+      cast_sinks_[1];
+  media_sink_service_->current_sinks_by_mdns_map_[params_[2].ip_endpoint] =
+      cast_sinks_[2];
+  // Cast sink 2, 3 from dial discovery
+  auto extra_data = cast_sinks_[2].cast_data();
+  extra_data.discovered_by_dial = true;
+  extra_data.model_name += " dial";
+
+  media_sink_service_->current_sinks_by_dial_map_[params_[2].ip_endpoint] =
+      cast_sinks_[2];
+  media_sink_service_->current_sinks_by_dial_map_[params_[3].ip_endpoint] =
+      cast_sinks_[3];
+
+  // Callback returns Cast sink 1, 2, 3
+  media_sink_service_->OnFetchCompleted();
+  EXPECT_TRUE(base::ContainsValue(sinks, cast_sinks_[1]));
+  EXPECT_TRUE(base::ContainsValue(sinks, cast_sinks_[2]));
+  EXPECT_TRUE(base::ContainsValue(sinks, cast_sinks_[3]));
+}
+
 }  // namespace media_router