Revert of Revert "[telemetry] Wire TsProxy through telemetry's network stack."

telemetry/telemetry/internal/browser/browser_unittest.py

 # Copyright 2012 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.
 
 import logging
 import os
 import shutil
 import tempfile
 import unittest
 
@@ skipping 165 matching lines @@
 def _GenerateBrowserProfile(number_of_tabs):
   """ Generate a browser profile which browser had |number_of_tabs| number of
   tabs opened before it was closed.
       Returns:
         profile_dir: the directory of profile.
   """
   profile_dir = tempfile.mkdtemp()
   options = options_for_unittests.GetCopy()
   options.browser_options.output_profile_path = profile_dir
   browser_to_create = browser_finder.FindBrowser(options)
-  with browser_to_create.Create(options) as browser:
-    browser.platform.SetHTTPServerDirectories(path.GetUnittestDataDir())
-    blank_file_path = os.path.join(path.GetUnittestDataDir(), 'blank.html')
-    blank_url = browser.platform.http_server.UrlOf(blank_file_path)
-    browser.foreground_tab.Navigate(blank_url)
-    browser.foreground_tab.WaitForDocumentReadyStateToBeComplete()
-    for _ in xrange(number_of_tabs - 1):
-      tab = browser.tabs.New()
-      tab.Navigate(blank_url)
-      tab.WaitForDocumentReadyStateToBeComplete()
-  return profile_dir
+  browser_to_create.platform.network_controller.InitializeIfNeeded()
+  try:
+    with browser_to_create.Create(options) as browser:
+      browser.platform.SetHTTPServerDirectories(path.GetUnittestDataDir())
+      blank_file_path = os.path.join(path.GetUnittestDataDir(), 'blank.html')
+      blank_url = browser.platform.http_server.UrlOf(blank_file_path)
+      browser.foreground_tab.Navigate(blank_url)
+      browser.foreground_tab.WaitForDocumentReadyStateToBeComplete()
+      for _ in xrange(number_of_tabs - 1):
+        tab = browser.tabs.New()
+        tab.Navigate(blank_url)
+        tab.WaitForDocumentReadyStateToBeComplete()
+    return profile_dir
+  finally:
+    browser_to_create.platform.network_controller.Close()
 
 
 class BrowserCreationTest(unittest.TestCase):
   def setUp(self):
     self.mock_browser_backend = mock.MagicMock()
     self.mock_platform_backend = mock.MagicMock()
 
   def testCleanedUpCalledWhenExceptionRaisedInBrowserCreation(self):
     self.mock_platform_backend.platform.FlushDnsCache.side_effect = (
         IntentionalException('Boom!'))
@@ skipping 19 matching lines @@
 
   @classmethod
   def setUpClass(cls):
     cls._number_of_tabs = 4
     cls._profile_dir = _GenerateBrowserProfile(cls._number_of_tabs)
     cls._options = options_for_unittests.GetCopy()
     cls._options.browser_options.AppendExtraBrowserArgs(
         ['--restore-last-session'])
     cls._options.browser_options.profile_dir = cls._profile_dir
     cls._browser_to_create = browser_finder.FindBrowser(cls._options)
+    cls._browser_to_create.platform.network_controller.InitializeIfNeeded()
 
   @decorators.Enabled('has tabs')
   @decorators.Disabled('chromeos', 'win', 'mac')
   # TODO(nednguyen): Enable this test on windowsn platform
   def testRestoreBrowserWithMultipleTabs(self):
     with self._browser_to_create.Create(self._options) as browser:
       # The number of tabs will be self._number_of_tabs + 1 as it includes the
       # old tabs and a new blank tab.
       expected_number_of_tabs = self._number_of_tabs + 1
       try:
         util.WaitFor(lambda: len(browser.tabs) == expected_number_of_tabs, 10)
       except:
         logging.error('Number of tabs is %s' % len(browser.tabs))
         raise
       self.assertEquals(expected_number_of_tabs, len(browser.tabs))
 
   @classmethod
   def tearDownClass(cls):
+    cls._browser_to_create.platform.network_controller.Close()
     shutil.rmtree(cls._profile_dir)
 
 
 class TestBrowserOperationDoNotLeakTempFiles(unittest.TestCase):
 
   @decorators.Enabled('win', 'mac', 'linux')
   @decorators.Isolated
   def testBrowserNotLeakingTempFiles(self):
     options = options_for_unittests.GetCopy()
     browser_to_create = browser_finder.FindBrowser(options)
     self.assertIsNotNone(browser_to_create)
     before_browser_run_temp_dir_content = os.listdir(tempfile.tempdir)
-    with browser_to_create.Create(options) as browser:
-      tab = browser.tabs.New()
-      tab.Navigate('about:blank')
-      self.assertEquals(2, tab.EvaluateJavaScript('1 + 1'))
-    after_browser_run_temp_dir_content = os.listdir(tempfile.tempdir)
-    self.assertEqual(before_browser_run_temp_dir_content,
-                     after_browser_run_temp_dir_content)
+    browser_to_create.platform.network_controller.InitializeIfNeeded()
+    try:
+      with browser_to_create.Create(options) as browser:
+        tab = browser.tabs.New()
+        tab.Navigate('about:blank')
+        self.assertEquals(2, tab.EvaluateJavaScript('1 + 1'))
+      after_browser_run_temp_dir_content = os.listdir(tempfile.tempdir)
+      self.assertEqual(before_browser_run_temp_dir_content,
+                       after_browser_run_temp_dir_content)
+    finally:
+      browser_to_create.platform.network_controller.Close()