#!/usr/bin/env python

#coding=utf-8  

import os

import platform

import tempfile

import shutil  

from PIL import Image  

PATH = lambda p: os.path.abspath(p)

TEMP_FILE = PATH(tempfile.gettempdir() + "/temp_screen.png")  

class Appium_Extend(object):

    def __init__(self, driver):

        self.driver = driver  

    def get_screenshot_by_element(self, element):

        #先截取整个屏幕，存储至系统临时目录下

        self.driver.get_screenshot_as_file(TEMP_FILE)  

        #获取元素bounds

        location = element.location

        size = element.size

        box = (location["x"], location["y"], location["x"] + size["width"], location["y"] + size["height"])  

        #截取图片

        image = Image.open(TEMP_FILE)

        newImage = image.crop(box)

        newImage.save(TEMP_FILE)  

        return self  

    def get_screenshot_by_custom_size(self, start_x, start_y, end_x, end_y):

        #自定义截取范围

        self.driver.get_screenshot_as_file(TEMP_FILE)

        box = (start_x, start_y, end_x, end_y)  

        image = Image.open(TEMP_FILE)

        newImage = image.crop(box)

        newImage.save(TEMP_FILE)  

        return self  

    def write_to_file( self, dirPath, imageName, form = "png"):

        #将截屏文件复制到指定目录下

        if not os.path.isdir(dirPath):

            os.makedirs(dirPath)

        shutil.copyfile(TEMP_FILE, PATH(dirPath + "/" + imageName + "." + form))  

    def load_image(self, image_path):

        #加载目标图片供对比用

        if os.path.isfile(image_path):

            load = Image.open(image_path)

            return load

        else:

            raise Exception("%s is not exist" %image_path)  

    #http://testerhome.com/topics/202

    def same_as(self, load_image, percent):

        #对比图片，percent值设为0，则100%相似时返回True，设置的值越大，相差越大

        import math

        import operator  

        image1 = Image.open(TEMP_FILE)

        image2 = load_image  

        histogram1 = image1.histogram()

        histogram2 = image2.histogram()  

        differ = math.sqrt(reduce(operator.add, list(map(lambda a,b: (a-b)**2, \

                                                         histogram1, histogram2)))/len(histogram1))

        if differ <= percent:

            return True

        else:

            return False  

#coding=utf-8  

import unittest

import os  

from extend import Appium_Extend

from appium import webdriver  

class Test(unittest.TestCase):

    #初始化环境

    def setUp(self):

        desired_caps = {}

        desired_caps["platformName"] = "Android"

        desired_caps["platformVersion"] = "4.3"

        desired_caps["deviceName"] = "788a6ab5"

        desired_caps["appPackage"] = "com.android.settings"

        desired_caps["appActivity"] = ".Settings"  

        self.driver = webdriver.Remote("http://127.0.0.1:4723/wd/hub", desired_caps)  

        self.extend = Appium_Extend(self.driver)  

        #回到主屏幕

        self.driver.press_keycode(3)  

    #退出测试

    def tearDown(self):

        self.driver.quit()  

    def test_get_screen_by_element(self):

        element = self.driver.find_element_by_id("com.android.deskclock:id/imageview")  

        self.extend.get_screenshot_by_element(element).write_to_file("f:\\screen", "image")

        self.assertTrue(os.path.isfile("f:\\screen\\image.png"))  

    def test_same_as(self):

        element = self.driver.find_element_by_id("com.android.deskclock:id/imageview")  

        load = self.extend.load_image("f:\\screen\\image.png")

        #要求百分百相似

        result = self.extend.get_screenshot_by_element(element).same_as(load, 0)

        self.assertTrue(result)  

if __name__ == "__main__":

    suite = unittest.TestSuite()

    suite.addTest(Test("test_get_screen_by_element"))

    suite.addTest(Test("test_same_as"))

    #执行测试

    unittest.TextTestRunner().run(suite)