autotest: make screenshot more robust.

Try to wake up screen if we can't find a valid crtc and try once
more.

BUG=chromium:607321
TEST=/usr/local/autotest/bin/screenshot.py /tmp/test.png

Change-Id: I669859bf652a8d178d240cbd20e1ac58d7007bf2
Reviewed-on: https://chromium-review.googlesource.com/341043
Commit-Ready: Ilja Friedel <ihf@chromium.org>
Tested-by: Ilja Friedel <ihf@chromium.org>
Reviewed-by: Stéphane Marchesin <marcheu@chromium.org>
Reviewed-by: Haixia Shi <hshi@chromium.org>
diff --git a/client/cros/graphics/drm.py b/client/cros/graphics/drm.py
index f5dfce7..637c228 100644
--- a/client/cros/graphics/drm.py
+++ b/client/cros/graphics/drm.py
@@ -14,6 +14,7 @@
 from ctypes import *
 import mmap
 import os
+import subprocess
 
 from PIL import Image
 
@@ -38,14 +39,12 @@
     _l = None
 
     def __repr__(self):
-        return "%s %d.%d.%d (%s) (%s)" % (
-            self.name,
-            self.version_major,
-            self.version_minor,
-            self.version_patchlevel,
-            self.desc,
-            self.date,
-        )
+        return "%s %d.%d.%d (%s) (%s)" % (self.name,
+                                          self.version_major,
+                                          self.version_minor,
+                                          self.version_patchlevel,
+                                          self.desc,
+                                          self.date,)
 
     def __del__(self):
         if self._l:
@@ -82,6 +81,17 @@
         if self._l:
             self._l.drmModeFreeResources(self)
 
+    def _wakeup_screen(self):
+        """
+        Send a synchronous dbus message to power on screen.
+        """
+        # Get and process reply to make this synchronous.
+        subprocess.check_output([
+            "dbus-send", "--type=method_call", "--system", "--print-reply",
+            "--dest=org.chromium.PowerManager", "/org/chromium/PowerManager",
+            "org.chromium.PowerManager.HandleUserActivity", "int32:0"
+        ])
+
     def getValidCrtc(self):
         for i in xrange(0, self.count_crtcs):
             crtc_id = self.crtcs[i]
@@ -90,19 +100,24 @@
                 return crtc
         return None
 
-    def getCrtc(self, crtc_id=None):
+    def getCrtc(self, crtc_id):
         """
         Obtain the CRTC at a given index.
 
         @param crtc_id: The CRTC to get.
         """
-        crtc = None
         if crtc_id:
-            crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents
-        else:
-            crtc = self.getValidCrtc()
-        crtc._fd = self._fd
-        crtc._l = self._l
+            return self._l.drmModeGetCrtc(self._fd, crtc_id).contents
+        return self.getValidCrtc()
+
+    def getCrtcRobust(self, crtc_id=None):
+        crtc = self.getCrtc(crtc_id)
+        if crtc is None:
+            self._wakeup_screen()
+            crtc = self.getCrtc(crtc_id)
+        if crtc is not None:
+            crtc._fd = self._fd
+            crtc._l = self._l
         return crtc
 
 
@@ -194,13 +209,11 @@
 
     def __repr__(self):
         s = "<Framebuffer (%dx%d (pitch %d bytes), %d bits/pixel, depth %d)"
-        vitals = s % (
-            self.width,
-            self.height,
-            self.pitch,
-            self.bpp,
-            self.depth,
-        )
+        vitals = s % (self.width,
+                      self.height,
+                      self.pitch,
+                      self.bpp,
+                      self.depth,)
         if self._map:
             tail = " (mapped)>"
         else:
@@ -232,8 +245,11 @@
         # mmap.mmap() has a totally different order of arguments in Python
         # compared to C; check the documentation before altering this
         # incantation.
-        self._map = mmap.mmap(self._fd, size, flags=mmap.MAP_SHARED,
-                              prot=mmap.PROT_READ, offset=mapDumb.offset)
+        self._map = mmap.mmap(self._fd,
+                              size,
+                              flags=mmap.MAP_SHARED,
+                              prot=mmap.PROT_READ,
+                              offset=mapDumb.offset)
 
     def unmap(self):
         """
@@ -382,6 +398,7 @@
 
 _drm = None
 
+
 def crtcScreenshot(crtc_id=None):
     """
     Take a screenshot, returning an image object.
@@ -399,9 +416,12 @@
                 break
 
     if _drm:
-        fb = _drm.resources().getCrtc(crtc_id).fb()
-        image = Image.new("RGB", (fb.width, fb.height))
-        pixels = _screenshot(image, fb)
-        return image
+        crtc = _drm.resources().getCrtcRobust(crtc_id)
+        if crtc is not None:
+            framebuffer = crtc.fb()
+            image = Image.new("RGB", (framebuffer.width, framebuffer.height))
+            pixels = _screenshot(image, framebuffer)
+            return image
 
-    raise RuntimeError("Couldn't screenshot with DRM devices")
+    raise RuntimeError(
+        "Unable to take screenshot. There may not be anything on the screen.")