Add extension policy tests

chrome/test/functional/policy.py

Index: chrome/test/functional/policy.py
===================================================================
--- chrome/test/functional/policy.py	(revision 126531)
+++ chrome/test/functional/policy.py	(working copy)
@@ -28,7 +28,7 @@
         self.GetPrefsInfo().Prefs(key), defaultval,
         msg='Expected preference value "%s" does not match actual value "%s".' %
         (defaultval, self.GetPrefsInfo().Prefs(key)))
-    self.assertRaises(pyauto.JSONInterfaceError,
+    self.assertRaises(pyauto_errors.JSONInterfaceError,
                       lambda: self.SetPrefs(key, newval))
 
   # TODO(frankf): Move tests dependending on this to plugins.py.
@@ -62,7 +62,7 @@
     try:
       ret = self.ExecuteJavascript('domAutomationController.send("done");')
       return ret == 'done'
-    except pyauto.JSONInterfaceError as e:
+    except pyauto_errors.JSONInterfaceError as e:
       if 'Javascript execution was blocked' == str(e):
         logging.debug('The previous failure was expected')
         return False
@@ -561,6 +561,135 @@
     self.assertRaises(pyauto.JSONInterfaceError,
                       lambda: self.OmniboxAcceptInput())
 
+  # Needed for extension tests
+  _GOOD_CRX_ID = 'ldnnhddmnhbkjipkidpdiheffobcpfmf'
+  _ADBLOCK_CRX_ID = 'dojnnbeimaimaojcialkkgajdnefpgcn'
+  _SCREEN_CAPTURE_CRX_ID = 'cpngackimfmofbokmjmljamhdncknpmg'
 
+  def _BuildCRXPath(self, crx_file_name):
+    """Returns the complete path to a crx_file in the data directory.
+
+    Args:
+      crx_file_name: The file name of the extension
+
+    Returns:
+      The full path to the crx in the data directory
+    """
+    return os.path.abspath(os.path.join(self.DataDir(), 'extensions',
+                                        crx_file_name))
+
+  def _AttemptExtensionInstallThatShouldFail(self, crx_file_name):
+    """Attempts to install an extension, raises an exception if installed.
+
+    Args:
+      crx_file_name: The file name of the extension
+    """
+    install_failed = True
+    try:
+      self.InstallExtension(self._BuildCRXPath(crx_file_name))
+      install_failed = False
+    except pyauto_errors.JSONInterfaceError as e:
+      self.assertEqual(e[0], 'Extension could not be installed',
+                       msg='The extension failed to install which is expected. '
+                       ' However it failed due to an unexpected reason: %s'
+                       % e[0])
+    self.assertTrue(install_failed, msg='The extension %s did not throw an '
+                    'exception when installation was attempted.  This most '
+                    'likely means it succeeded, which it should not.')
+
+  def _AttemptExtensionInstallThatShouldPass(self, crx_file_name):
+    """Attempts to install an extension, raises an exception if not installed.
+
+    Args:
+      crx_file_name: The file name of the extension
+    """
+    ext_id = self.InstallExtension(self._BuildCRXPath(crx_file_name))
+    self.assertTrue(self._CheckForExtensionByID(ext_id),
+                    msg='The %s extension was not install even though '
+                    'it should have been.' % crx_file_name)
+
+  def _CheckForExtensionByID(self, extension_id):
+    """Returns if an extension is installed.
+
+    Args:
+      extension_id: The id of the extension
+
+    Returns:
+      True if the extension is installed; False otherwise
+    """
+    all_extensions = [extension['id'] for extension in self.GetExtensionsInfo()]
+    return extension_id in all_extensions
+
+  def _RemoveTestingExtensions(self):
+    """Temporary method that cleans state for extension test.
+
+    Currently the tear down method for policy_base does not clear the user
+    state.  See crosbug.com/27227.
+    """
+    if self._CheckForExtensionByID(self._SCREEN_CAPTURE_CRX_ID):
+      self.UninstallExtensionById(self._SCREEN_CAPTURE_CRX_ID)
+    if self._CheckForExtensionByID(self._GOOD_CRX_ID):
+      self.UninstallExtensionById(self._GOOD_CRX_ID)
+      self.NavigateToURL('chrome:extensions')
+    if self._CheckForExtensionByID(self._ADBLOCK_CRX_ID):
+      self.UninstallExtensionById(self._ADBLOCK_CRX_ID)
+    # There is an issue where if you uninstall and reinstall and extension
+    # quickly with self.InstallExtension, the reinstall fails.  This is a hack
+    # to fix it.  Bug coming soon.
+    self.NavigateToURL('chrome:extensions')
+
+  def testExtensionInstallPopulatedBlacklist(self):
+    """Verify blacklisted extensions cannot be installed."""
+    # TODO(krisr): Remove this when crosbug.com/27227 is fixed.
+    self._RemoveTestingExtensions()
+    # Blacklist good.crx
+    self.SetPolicies({
+      'ExtensionInstallBlacklist': [self._GOOD_CRX_ID]
+    })
+    self._AttemptExtensionInstallThatShouldFail('good.crx')
+    # Check adblock is installed.
+    self._AttemptExtensionInstallThatShouldPass('adblock.crx')
+
+  def testExtensionInstallFailWithGlobalBlacklist(self):
+    """Verify no extensions can be installed when all are blacklisted."""
+    # TODO(krisr): Remove this when crosbug.com/27227 is fixed.
+    self._RemoveTestingExtensions()
+    # Block installs of all extensions
+    self.SetPolicies({
+      'ExtensionInstallBlacklist': ['*']
+    })
+    self._AttemptExtensionInstallThatShouldFail('good.crx')
+    self._AttemptExtensionInstallThatShouldFail('adblock.crx')
+
+  def testExtensionInstallWithGlobalBlacklistAndWhitelistedExtension(self):
+    """Verify whitelisted extension is installed when all are blacklisted."""
+    # TODO(krisr): Remove this when crosbug.com/27227 is fixed.
+    self._RemoveTestingExtensions()
+    # Block installs of all extensions, but whitelist adblock.crx
+    self.SetPolicies({
+      'ExtensionInstallBlacklist': ['*'],
+      'ExtensionInstallWhitelist': [self._ADBLOCK_CRX_ID]
+    })
+    self._AttemptExtensionInstallThatShouldFail('good.crx')
+    self._AttemptExtensionInstallThatShouldPass('adblock.crx')
+
+  # TODO(krisr): Enable this test once we figure out why it isn't downloading
+  # the extension, crbug.com/118123.
+  def testExtensionInstallFromForceList(self):
+    """Verify force install extensions are installed."""
+    # TODO(krisr): Remove this when crosbug.com/27227 is fixed.
+    self._RemoveTestingExtensions()
+    # Force an extension download from the webstore.
+    self.SetPolicies({
+      'ExtensionInstallForcelist': [self._SCREEN_CAPTURE_CRX_ID],
+    })
+    # Give the system 30 seconds to go get this extension.  We are not sure how
+    # long it will take the policy to take affect and download the extension.
+    self.assertTrue(self.WaitUntil(lambda:
+      self._CheckForExtensionByID(self._SCREEN_CAPTURE_CRX_ID),
+      expect_retval=True),
+      msg='The force install extension was never installed.')
+
+
 if __name__ == '__main__':
   pyauto_functional.Main()