Define screenshot methods in this lib (#98)

* tweak tests for screenshots

* extract screenshots

* fix rubocop

* fix wrong tests

* tweak docstring

Author: KazuCocoa

CHANGELOG.md

@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 ### Enhancements
 - Add a `ImageElement` to handle images as elements by `matchTemplate`
+- [Internal] Define screenshot methods in appium_lib_core instead of Selenium's one
 
 ### Bug fixes
 

lib/appium_lib_core/common/base/bridge/mjsonwp.rb

@@ -6,6 +6,74 @@ class MJSONWP < ::Selenium::WebDriver::Remote::OSS::Bridge
           def commands(command)
             ::Appium::Core::Commands::MJSONWP::COMMANDS[command]
           end
+
+          #
+          # @return [::Appium::Core::ImageElement|nil]
+          # @raise [::Selenium::WebDriver::Error::TimeOutError|::Selenium::WebDriver::Error::WebDriverError]
+          #
+          def find_element_by_image(full_image:, partial_image:, match_threshold: nil, visualize: false)
+            options = {}
+            options[:threshold] = match_threshold unless match_threshold.nil?
+            options[:visualize] = visualize
+
+            params = {}
+            params[:mode] = :matchTemplate
+            params[:firstImage] = full_image
+            params[:secondImage] = partial_image
+            params[:options] = options if options
+
+            result = execute(:compare_images, {}, params)
+            rect = result['rect']
+
+            if rect
+              return ::Appium::Core::ImageElement.new(self,
+                                                      rect['x'],
+                                                      rect['y'],
+                                                      rect['width'],
+                                                      rect['height'],
+                                                      result['visualization'])
+            end
+            nil
+          end
+
+          #
+          # @return [[]|[::Appium::Core::ImageElement]]
+          # @raise [::Selenium::WebDriver::Error::TimeOutError|::Selenium::WebDriver::Error::WebDriverError]
+          #
+          def find_elements_by_image(full_image:, partial_images:, match_threshold: nil, visualize: false)
+            options = {}
+            options[:threshold] = match_threshold unless match_threshold.nil?
+            options[:visualize] = visualize
+
+            params = {}
+            params[:mode] = :matchTemplate
+            params[:firstImage] = full_image
+            params[:options] = options if options
+
+            partial_images.each_with_object([]) do |partial_image, acc|
+              params[:secondImage] = partial_image
+
+              begin
+                result = execute(:compare_images, {}, params)
+                rect = result['rect']
+
+                if result['rect']
+                  acc.push ::Appium::Core::ImageElement.new(self,
+                                                            rect['x'],
+                                                            rect['y'],
+                                                            rect['width'],
+                                                            rect['height'],
+                                                            result['visualization'])
+                end
+              rescue ::Selenium::WebDriver::Error::WebDriverError => e
+                acc if e.message.include?('Cannot find any occurrences')
+              end
+            end
+          end
+
+          def take_element_screenshot(element)
+            execute :take_element_screenshot, id: element.ref
+          end
         end # class MJSONWP
       end # class Bridge
     end # class Base

lib/appium_lib_core/common/base/bridge/w3c.rb

@@ -204,6 +204,10 @@ def log(type)
             end
           end
 
+          def take_element_screenshot(element)
+            execute :take_element_screenshot, id: element.ref
+          end
+
           private
 
           # Don't convert locators for Appium Client

lib/appium_lib_core/common/base/driver.rb

@@ -1,18 +1,19 @@
 require 'base64'
 require_relative 'search_context'
+require_relative 'screenshot'
 
 module Appium
   module Core
     class Base
       class Driver < ::Selenium::WebDriver::Driver
         include ::Selenium::WebDriver::DriverExtensions::UploadsFiles
-        include ::Selenium::WebDriver::DriverExtensions::TakesScreenshot
         include ::Selenium::WebDriver::DriverExtensions::HasSessionId
         include ::Selenium::WebDriver::DriverExtensions::Rotatable
         include ::Selenium::WebDriver::DriverExtensions::HasRemoteStatus
         include ::Selenium::WebDriver::DriverExtensions::HasWebStorage
 
         include ::Appium::Core::Base::SearchContext
+        include ::Appium::Core::Base::TakeScreenshot
 
         def initialize(opts = {})
           listener = opts.delete(:listener)

lib/appium_lib_core/common/base/screenshot.rb

@@ -0,0 +1,83 @@
+module Appium
+  module Core
+    class Base
+      module TakeScreenshot
+        #
+        # Save a PNG screenshot to the given path
+        #
+        # @api public
+        #
+        def save_screenshot(png_path)
+          extension = File.extname(png_path).downcase
+          if extension != '.png'
+            ::Appium::Logger.warn 'name used for saved screenshot does not match file type. '\
+                                  'It should end with .png extension'
+          end
+          File.open(png_path, 'wb') { |f| f << screenshot_as(:png) }
+        end
+
+        #
+        # Return a PNG screenshot in the given format as a string
+        #
+        # @param [:base64, :png] format
+        # @return String screenshot
+        #
+        # @example
+        #
+        #     @@driver.screenshot_as :base64 #=> "iVBORw0KGgoAAAANSUhEUgAABDgAAAB+CAIAAABOPDa6AAAAAX"
+        #
+        # @api public
+        def screenshot_as(format)
+          case format
+          when :base64
+            bridge.screenshot
+          when :png
+            bridge.screenshot.unpack('m')[0]
+          else
+            raise Core::Error::UnsupportedOperationError, "unsupported format: #{format.inspect}"
+          end
+        end
+
+        # @param [Selenium::WebDriver::Element] element A element you'd like to take screenshot.
+        # @param [String] png_path A path to save the screenshot
+        # @return [File] Path to the screenshot.
+        #
+        # @example
+        #
+        #   @driver.save_element_screenshot(element, "fine_name.png")
+        #
+        def save_element_screenshot(element, png_path)
+          extension = File.extname(png_path).downcase
+          if extension != '.png'
+            ::Appium::Logger.warn 'name used for saved screenshot does not match file type. '\
+                                'It should end with .png extension'
+          end
+          File.open(png_path, 'wb') { |f| f << element_screenshot_as(element, :png) }
+        end
+        # backward compatibility
+        alias take_element_screenshot save_element_screenshot
+
+        #
+        # Return a PNG screenshot in the given format as a string
+        #
+        # @param [:base64, :png] format
+        # @return String screenshot
+        #
+        # @example
+        #
+        #     @@driver.element_screenshot_as element, :base64 #=> "iVBORw0KGgoAAAANSUhEUgAABDgAAAB+CAIAAABOPDa6AAAAAX"
+        #
+        def element_screenshot_as(element, format)
+          case format
+          when :base64
+            bridge.take_element_screenshot(element)
+          when :png
+            bridge.take_element_screenshot(element).unpack('m')[0]
+          else
+            raise Core::Error::UnsupportedOperationError, "unsupported format: #{format.inspect}"
+          end
+        end
+      end
+    end
+  end
+end

lib/appium_lib_core/common/error.rb

@@ -11,6 +11,8 @@ class CapabilityStructureError < CoreError; end
       class NotSupportedAppiumServer < CoreError; end
       class NoSuchElementError < CoreError; end
 
+      class UnsupportedOperationError < CoreError; end
+
       # Server side error
       class ServerError; end
     end

lib/appium_lib_core/device.rb

@@ -477,16 +477,6 @@ module Device
       #   @driver.switch_to_default_context
       #
 
-      # @!method take_element_screenshot(element, png_path)
-      # @param [Selenium::WebDriver::Element] element A element you'd like to take screenshot.
-      # @param [String] png_path A path to save the screenshot
-      # @return [File] Path to the screenshot.
-      #
-      # @example
-      #
-      #   @driver.take_element_screenshot(element, "fine_name.png")
-      #
-
       # @!method stop_recording_screen(remote_path: nil, user: nil, pass: nil, method: 'PUT')
       # @param [String] remote_path: The path to the remote location, where the resulting video should be uploaded.
       #                             The following protocols are supported: http/https, ftp.
@@ -536,19 +526,6 @@ def device_time(format = nil)
             end
           end
 
-          add_endpoint_method(:take_element_screenshot) do
-            def take_element_screenshot(element, png_path)
-              result = execute :take_element_screenshot, id: element.ref
-
-              extension = File.extname(png_path).downcase
-              if extension != '.png'
-                ::Appium::Logger.warn 'name used for saved screenshot does not match file type. '\
-                                'It should end with .png extension'
-              end
-              File.open(png_path, 'wb') { |f| f << result.unpack('m')[0] }
-            end
-          end
-
           add_endpoint_method(:save_viewport_screenshot) do
             def save_viewport_screenshot(png_path)
               extension = File.extname(png_path).downcase
@@ -574,6 +551,10 @@ def save_viewport_screenshot(png_path)
           ScreenRecord.add_methods
           ImageComparison.add_methods
           AppState.add_methods
+
+          # Compatibility for appium_lib
+          # TODO: Will remove
+          delegate_from_appium_driver :take_element_screenshot
         end
 
         # def extended

test/unit/android/device/mjsonwp/definition_test.rb

@@ -19,12 +19,15 @@ def test_delegate_driver_method
             assert @driver.respond_to? :launch_app
           end
 
-          def test_delegate_from_appium_driver
-            assert @core.send(:delegated_target_for_test).respond_to? :launch_app
+          def delegate_from_appium_driver(key)
+            assert @core.send(:delegated_target_for_test).respond_to? key
           end
 
           def parameterized_method_defined_check(array)
-            array.each { |v| assert ::Appium::Core::Base::Bridge::MJSONWP.method_defined?(v) }
+            array.each do |v|
+              assert ::Appium::Core::Base::Bridge::MJSONWP.method_defined?(v)
+              delegate_from_appium_driver(v)
+            end
           end
 
           def test_with_arg_definitions

test/unit/android/device/w3c/definition_test.rb

@@ -18,12 +18,15 @@ def test_delegate_driver_method
             assert @driver.respond_to? :launch_app
           end
 
-          def test_delegate_from_appium_driver
-            assert @core.send(:delegated_target_for_test).respond_to? :launch_app
+          def delegate_from_appium_driver(key)
+            assert @core.send(:delegated_target_for_test).respond_to? key
           end
 
           def parameterized_method_defined_check(array)
-            array.each { |v| assert ::Appium::Core::Base::Bridge::W3C.method_defined?(v) }
+            array.each do |v|
+              assert ::Appium::Core::Base::Bridge::W3C.method_defined?(v)
+              delegate_from_appium_driver(v)
+            end
           end
 
           def test_with_arg_definitions

test/unit/ios/device/mjsonwp/definition_test.rb

@@ -18,12 +18,15 @@ def test_delegate_driver_method
             assert @driver.respond_to? :launch_app
           end
 
-          def test_delegate_from_appium_driver
-            assert @core.send(:delegated_target_for_test).respond_to? :launch_app
+          def delegate_from_appium_driver(key)
+            assert @core.send(:delegated_target_for_test).respond_to? key
           end
 
           def parameterized_method_defined_check(array)
-            array.each { |v| assert ::Appium::Core::Base::Bridge::MJSONWP.method_defined?(v) }
+            array.each do |v|
+              assert ::Appium::Core::Base::Bridge::MJSONWP.method_defined?(v)
+              delegate_from_appium_driver(v)
+            end
           end
 
           def test_with_arg_definitions

test/unit/ios/device/w3c/definition_test.rb

@@ -22,8 +22,15 @@ def test_delegate_from_appium_driver
             assert @core.send(:delegated_target_for_test).respond_to? :launch_app
           end
 
+          def delegate_from_appium_driver(key)
+            assert @core.send(:delegated_target_for_test).respond_to? key
+          end
+
           def parameterized_method_defined_check(array)
-            array.each { |v| assert ::Appium::Core::Base::Bridge::W3C.method_defined?(v) }
+            array.each do |v|
+              assert ::Appium::Core::Base::Bridge::W3C.method_defined?(v)
+              delegate_from_appium_driver(v)
+            end
           end
 
           def test_with_arg_definitions