Send screenshots back to the client for debugging

chrome/test/webdriver/chromedriver_tests.py

 #!/usr/bin/python
 
 # Copyright (c) 2011 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.
 
 """Tests for ChromeDriver.
 
 If your test is testing a specific part of the WebDriver API, consider adding
 it to the appropriate place in the WebDriver tree instead.
 """
 
+import hashlib
+import os
 import platform
-import os
 import sys
 import unittest
+import urllib
 import urllib2
 import urlparse
 
 from chromedriver_launcher import ChromeDriverLauncher
 import chromedriver_paths
 from gtest_text_test_runner import GTestTextTestRunner
 
 sys.path += [chromedriver_paths.SRC_THIRD_PARTY]
 sys.path += [chromedriver_paths.PYTHON_BINDINGS]
 
@@ skipping 49 matching lines @@
 
 class BasicTest(unittest.TestCase):
   """Basic ChromeDriver tests."""
 
   def setUp(self):
     self._launcher = ChromeDriverLauncher()
 
   def tearDown(self):
     self._launcher.Kill()
 
+  @staticmethod
+  def DataDir():
+    """Returns the path to the data dir chrome/test/data."""
+    return os.path.normpath(
+      os.path.join(os.path.dirname(__file__), os.pardir, "data"))
+
+  @staticmethod
+  def GetFileURLForPath(path):
+    """Get file:// url for the given path.
+    Also quotes the url using urllib.quote().
+    """
+    abs_path = os.path.abspath(path)
+    if sys.platform == 'win32':
+      # Don't quote the ':' in drive letter ( say, C: ) on win.
+      # Also, replace '\' with '/' as expected in a file:/// url.
+      drive, rest = os.path.splitdrive(abs_path)
+      quoted_path = drive.upper() + urllib.quote((rest.replace('\\', '/')))
+      return 'file:///' + quoted_path
+    else:
+      quoted_path = urllib.quote(abs_path)
+      return 'file://' + quoted_path
+
   def testShouldReturn404WhenSentAnUnknownCommandURL(self):
     request_url = self._launcher.GetURL() + '/foo'
     try:
       SendRequest(request_url, method='GET')
       self.fail('Should have raised a urllib.HTTPError for returned 404')
     except urllib2.HTTPError, expected:
       self.assertEquals(404, expected.code)
 
   def testShouldReturnHTTP405WhenSendingANonPostToTheSessionURL(self):
     request_url = self._launcher.GetURL() + '/session'
@@ skipping 45 matching lines @@
     self.assertNotEqual(cookie_dict, None)
     self.assertEqual(cookie_dict["value"], "this is a test");
 
   def testDeleteCookie(self):
     self.testAddCookie();
     self._driver.delete_cookie("chromedriver_cookie_test")
     cookie_dict = self._driver.get_cookie("chromedriver_cookie_test")
     self.assertEqual(cookie_dict, None)
 
 
+class ScreenshotTest(unittest.TestCase):
+  """Tests to verify screenshot retrieval"""
+
+  SEARCH = "http://www.google.com/webhp?hl=en"
+  NEWS = "http://www.google.com/news"
+  REDBOX = "automation_proxy_snapshot/set_size.html"
+
+  def setUp(self):
+    self._launcher = ChromeDriverLauncher()
+    self._driver = WebDriver(self._launcher.GetURL(), {})
+
+  def tearDown(self):
+    self._driver.quit()
+    self._launcher.Kill()
+
+  def testScreenCaptureAgainstReference(self):
+    # Create a red square of 2000x2000 pixels.
+    url = BasicTest.GetFileURLForPath(os.path.join(BasicTest.DataDir(),
+                                                   self.REDBOX))
+    url = url + "?2000,2000"
+    self._driver.get(url)
+    s = self._driver.get_screenshot_as_base64();
+    h = hashlib.md5(s).hexdigest()
+    # Compare the PNG created to the reference hash.
+    self.assertTrue(h, '12c0ade27e3875da3d8866f52d2fa84f')
+
 class SessionTest(unittest.TestCase):
   """Tests dealing with WebDriver sessions."""
 
   def setUp(self):
     self._launcher = ChromeDriverLauncher()
 
   def tearDown(self):
     self._launcher.Kill()
 
   def testCreatingSessionShouldRedirectToCorrectURL(self):
@@ skipping 127 matching lines @@
     self.assertEquals('', url_parts[0])
     self.assertEquals('wd', url_parts[1])
     self.assertEquals('hub', url_parts[2])
     self.assertEquals('session', url_parts[3])
     self.assertEquals(data['sessionId'], url_parts[4])
 
 
 if __name__ == '__main__':
   unittest.main(module='chromedriver_tests',
                 testRunner=GTestTextTestRunner(verbosity=1))