Support device.screenshot() API

Change-Id: I2b83a2e2e9698342fcf193a9fc129640d873c6b0

16 files changed, 109 insertions, 17 deletions

diff --git a/CHANGES.txt b/CHANGES.txt
index 5655b55..9da8bc4 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -39,3 +39,7 @@ Version 0.3.1   22 Sep 2016
 Version 0.3.2   07 Oct 2016
 ---------------------------
 - Support artik and more standalone device types
+
+Version 0.3.3   11 Oct 2016
+---------------------------
+- Support device.screenshot() API
diff --git a/README.md b/README.md
index 179460a..a5124e4 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,7 @@ Buliding & installing
 
    $ cd ..
    
-   $ sudo dpkg -i litmus_0.3.1-1_amd64.deb
+   $ sudo dpkg -i litmus_0.3.3-1_amd64.deb
 
 
 Getting started
diff --git a/debian/changelog b/debian/changelog
index 48b9aa1..63cb7d0 100644
--- a/debian/changelog
+++ b/debian/changelog
@@ -1,3 +1,9 @@
+litmus (0.3.3-1) unstable; urgency=low
+
+  * Support device.screenshot API
+
+ -- Donghoon Shin <dhs.shin@samsung.com>  Tue, 11 Oct 2016 12:41:00 +0900
+
 litmus (0.3.2-1) unstable; urgency=low
 
   * Support artik and more standalone device types
diff --git a/debian/control b/debian/control
index ab23163..2d2d43c 100644
--- a/debian/control
+++ b/debian/control
@@ -15,9 +15,11 @@ Depends: ${python3:Depends}, ${misc:Depends},
  python3-yaml (>= 3.10),
  python3-requests (>= 2.2.1),
  python3-bs4 (>= 4.2.1),
+ python3-pil (>= 2.3.0),
  python3-fasteners (>= 0.12),
  git (>= 1.9),
  lthor (>= 2.0),
+ imagemagick (>= 6.7.7),
  clewarecontrol (>= 4.1),
  smartpower (>= 0.1),
  heimdall-flash (>= 1.4.1-2),
diff --git a/litmus/__init__.py b/litmus/__init__.py
index 131d9e3..9a396a8 100644
--- a/litmus/__init__.py
+++ b/litmus/__init__.py
@@ -14,7 +14,7 @@
 # limitations under the License.
 import os
 
-__version__ = '0.3.2'
+__version__ = '0.3.3'
 _homedir_ = os.path.expanduser('~')
 _confdir_ = os.path.join(_homedir_, '.litmus')
 _duts_ = os.path.join(_confdir_, 'topology')
diff --git a/litmus/device/device.py b/litmus/device/device.py
index 981f7d5..2908898 100644
--- a/litmus/device/device.py
+++ b/litmus/device/device.py
@@ -17,7 +17,9 @@ import os
 import time
 import serial
 import logging
+import tempfile
 import fasteners
+from PIL import Image
 from threading import Thread, Lock
 from litmus.core.util import call, check_output
 from litmus.core.util import convert_single_item_to_list
@@ -52,6 +54,8 @@ class device(object):
     _max_attempt_boot_retry = 3
     _boot_timeout = 50.0
     _path_for_locks = _path_for_locks_
+    _screen_width = 1920
+    _screen_height = 1080
 
     _cutter = None
     _uart = None
@@ -349,7 +353,86 @@ class device(object):
             elif isinstance(l['args'], tuple):
                 l['func'](self, *l['args'])
 
+    def screenshot(self, filename):
+        """
+        Take a screenshot (png format)
+
+        :param str filename: screenshot file name
+
+        Example:
+            >>> dut.screenshot('screenshot.png')
+
+        """
+        logging.debug('==== Take a screenshot: {}, width: {}, height: {} ===='
+                      .format(filename,
+                              self._screen_width,
+                              self._screen_height))
+        dt = self._get_display_server_type()
+        if dt == 'X11':
+            self._screenshot_x11(filename)
+        elif dt == 'WAYLAND':
+            self._screenshot_wayland(filename)
+
     # private methods.
+    def _get_display_server_type(self):
+        """docstring for get_display_server_type"""
+        res = self.run_cmd('ls /usr/lib', timeout=10)
+        if find_pattern('.*libX11.*', res):
+            return 'X11'
+        else:
+            return 'WAYLAND'
+
+    def _screenshot_x11(self, filename):
+        """docstring for _screenshot_x11"""
+        # take a screenshot using xwd
+        cmd = 'xwd -root -out /tmp/{}.xwd'.format(filename)
+        self.run_cmd(cmd, timeout=20)
+
+        # pull xwd file
+        self.pull_file('/tmp/{}.xwd'.format(filename), os.curdir, timeout=20)
+
+        # convert xwd to png and resize it
+        call(['convert', '{0}.xwd'.format(filename), '-resize',
+              '{0}x{1}'.format(self._screen_width, self._screen_height),
+              filename], timeout=20)
+        call(['rm', '{}.xwd'.format(filename)], timeout=20)
+
+    def _screenshot_wayland(self, filename):
+        """docstring for _screenshot_wayland"""
+        # Find all viewable window id
+        p_winid = '.*(0x[a-zA-Z0-9]{8})\s+\d+\s+\d+\s+\d+\s+\d+\s+(\d+)\s+(\d+).*\sViewable.*'
+        winids = find_all_pattern(p_winid,
+                                  self.run_cmd('enlightenment_info -topvwins',
+                                               timeout=20))
+        if winids:
+            # Dump windows
+            outs = self.run_cmd('enlightenment_info -dump_topvwins',
+                                timeout=20)
+            dirn = find_pattern('directory:\s(.*)',
+                                outs,
+                                groupindex=1).rstrip()
+
+            # Create tempdir and pull dump files
+            tmpdir = tempfile.mkdtemp()
+            self.pull_file(dirn, tmpdir, timeout=20)
+
+            # If dump does not exist then remove winid from list
+            winids = [winid for winid in winids
+                      if os.path.exists(os.path.join(tmpdir, winid[0]+'.png'))]
+
+            # Base image
+            bg = Image.new('RGB', (self._screen_width, self._screen_height))
+            # Merge images
+            for winid in reversed(winids):
+                try:
+                    fg = Image.open(os.path.join(tmpdir, winid[0]+'.png'))
+                    bg.paste(fg, (int(winid[1]), int(winid[2])), fg)
+                except FileNotFoundError:
+                    pass
+            # Save merged image
+            bg.save(filename)
+            # Remove tempdir
+            call(['rm', '-rf', tmpdir], timeout=10)
 
     def _flush_uart_buffer(self):
         """docstring for flush_uart_buffer"""
diff --git a/litmus/device/deviceartik5.py b/litmus/device/deviceartik5.py
index bc4323e..15129af 100644
--- a/litmus/device/deviceartik5.py
+++ b/litmus/device/deviceartik5.py
@@ -19,3 +19,5 @@ from litmus.device.deviceartik10 import deviceartik10
 class deviceartik5(deviceartik10):
     """docstring for device"""
     _pattern_bootprompt = r'ARITK.*# .*'
+    _screen_width = 480
+    _screen_height = 800
diff --git a/litmus/device/devicestandalone_m0.py b/litmus/device/devicestandalone_m0.py
index 41cc556..6306d89 100644
--- a/litmus/device/devicestandalone_m0.py
+++ b/litmus/device/devicestandalone_m0.py
@@ -18,5 +18,5 @@ from litmus.device.devicestandalone import devicestandalone
 
 class devicestandalone_m0(devicestandalone):
     """docstring for device"""
-    pass
-
+    _screen_width = 720
+    _screen_height = 1280
diff --git a/litmus/device/devicestandalone_tm1.py b/litmus/device/devicestandalone_tm1.py
index 0cada07..b6f9a7e 100644
--- a/litmus/device/devicestandalone_tm1.py
+++ b/litmus/device/devicestandalone_tm1.py
@@ -18,5 +18,5 @@ from litmus.device.devicestandalone import devicestandalone
 
 class devicestandalone_tm1(devicestandalone):
     """docstring for device"""
-    pass
-
+    _screen_width = 720
+    _screen_height = 1280
diff --git a/litmus/device/devicestandalone_tm2.py b/litmus/device/devicestandalone_tm2.py
index 79cb68b..208a7d5 100644
--- a/litmus/device/devicestandalone_tm2.py
+++ b/litmus/device/devicestandalone_tm2.py
@@ -18,5 +18,5 @@ from litmus.device.devicestandalone import devicestandalone
 
 class devicestandalone_tm2(devicestandalone):
     """docstring for device"""
-    pass
-
+    _screen_width = 1440
+    _screen_height = 2560
diff --git a/litmus/device/devicestandalone_tw1.py b/litmus/device/devicestandalone_tw1.py
index fe459e8..774a250 100644
--- a/litmus/device/devicestandalone_tw1.py
+++ b/litmus/device/devicestandalone_tw1.py
@@ -18,5 +18,5 @@ from litmus.device.devicestandalone import devicestandalone
 
 class devicestandalone_tw1(devicestandalone):
     """docstring for device"""
-    pass
-
+    _screen_width = 360
+    _screen_height = 360
diff --git a/litmus/device/devicestandalone_u3.py b/litmus/device/devicestandalone_u3.py
index 9517b1e..eae7668 100644
--- a/litmus/device/devicestandalone_u3.py
+++ b/litmus/device/devicestandalone_u3.py
@@ -19,4 +19,3 @@ from litmus.device.devicestandalone import devicestandalone
 class devicestandalone_u3(devicestandalone):
     """docstring for device"""
     pass
-
diff --git a/litmus/device/devicestandalone_xu3.py b/litmus/device/devicestandalone_xu3.py
index 08bc420..44d8206 100644
--- a/litmus/device/devicestandalone_xu3.py
+++ b/litmus/device/devicestandalone_xu3.py
@@ -19,4 +19,3 @@ from litmus.device.devicestandalone import devicestandalone
 class devicestandalone_xu3(devicestandalone):
     """docstring for device"""
     pass
-
diff --git a/litmus/device/devicestandalone_z1.py b/litmus/device/devicestandalone_z1.py
index 8a85f76..a4e61f5 100644
--- a/litmus/device/devicestandalone_z1.py
+++ b/litmus/device/devicestandalone_z1.py
@@ -18,5 +18,5 @@ from litmus.device.devicestandalone import devicestandalone
 
 class devicestandalone_z1(devicestandalone):
     """docstring for device"""
-    pass
-
+    _screen_width = 480
+    _screen_height = 800
diff --git a/litmus/device/deviceu3.py b/litmus/device/deviceu3.py
index 9c5d399..a8687d3 100644
--- a/litmus/device/deviceu3.py
+++ b/litmus/device/deviceu3.py
@@ -19,4 +19,3 @@ from litmus.device.device import device
 class deviceu3(device):
     """docstring for device"""
     pass
-
diff --git a/litmus/helper/helper.py b/litmus/helper/helper.py
index 26a8cf2..0b4eac3 100644
--- a/litmus/helper/helper.py
+++ b/litmus/helper/helper.py
@@ -18,7 +18,6 @@ import sys
 import time
 import logging
 import requests
-import litmus
 import urllib.parse
 from bs4 import BeautifulSoup
 from litmus.core.util import find_pattern, find_all_pattern
@@ -173,7 +172,6 @@ def install_plugin(dut, script, waiting=5, timeout=180):
     dut.off()
 
 
-import os
 import tempfile
 import shutil
 from subprocess import DEVNULL