Update caching to be on BaseConnection

Signed-off-by: Jared Quick <jquick@chef.io>

lib/train/plugins/base_connection.rb

@@ -22,9 +22,19 @@ class BaseConnection
     def initialize(options = nil)
       @options = options || {}
       @logger = @options.delete(:logger) || Logger.new(STDOUT)
+      @cacher = CacheConnection.new(self)
       Train::Platforms::Detect::Specifications::OS.load
     end
 
+    def enable_cache(type)
+      @cacher.cache_enabled[type.to_sym] = true
+    end
+
+    def disable_cache(type)
+      @cacher.cache_enabled[type.to_sym] = false
+      @cacher.clear_cache(type.to_sym)
+    end
+
     # Closes the session connection, if it is still active.
     def close
       # this method may be left unimplemented if that is applicable
@@ -48,31 +58,44 @@ def local?
       false
     end
 
-    # Execute a command using this connection.
-    #
-    # @param command [String] command string to execute
-    # @return [CommandResult] contains the result of running the command
-    def run_command(_command)
-      fail Train::ClientError, "#{self.class} does not implement #run_command()"
-    end
-
     # Get information on the operating system which this transport connects to.
     #
     # @return [Platform] system information
     def platform
       @platform ||= Train::Platforms::Detect.scan(self)
     end
-
     # we need to keep os as a method for backwards compatibility with inspec
     alias os platform
 
+    # Execute a command using this connection.
+    #
+    # @param command [String] command string to execute
+    # @return [CommandResult] contains the result of running the command
+    def run_command_via_connection(_command)
+      fail Train::ClientError, "#{self.class} does not implement #run_command_via_connection()"
+    end
+
+    # run command with optional caching
+    def run_command(command)
+      return @cacher.run_command(command) if @cacher.cache_enabled[:command] == true
+
+      run_command_via_connection(command)
+    end
+
     # Interact with files on the target. Read, write, and get metadata
     # from files via the transport.
     #
     # @param [String] path which is being inspected
     # @return [FileCommon] file object that allows for interaction
-    def file(_path, *_args)
-      fail Train::ClientError, "#{self.class} does not implement #file(...)"
+    def file_via_connection(_path, *_args)
+      fail Train::ClientError, "#{self.class} does not implement #file_via_connection(...)"
+    end
+
+    # file with optional caching
+    def file(path, *args)
+      return @cacher.file(path, *args) if @cacher.cache_enabled[:file] == true
+
+      file_via_connection(path, *args)
     end
 
     # Builds a LoginCommand which can be used to open an interactive

lib/train/plugins/cache_connection.rb

@@ -1,37 +1,44 @@
 # encoding: utf-8
 
 class Train::Plugins::Transport
-  class CacheConnection < BaseConnection
+  class CacheConnection
+    attr_accessor :cache_enabled
     # Create a new CacheConnection instance. This instance will cache
     # file and command operations on the underline connection.
     #
     # @param connection [Connection] connection object
     def initialize(connection)
       @connection = connection
-      @file_cache = {}
-      @cmd_cache = {}
+      @cache = {}
 
-      @connection.class.instance_methods(false).each do |m|
-        next if %i(run_command file).include?(m)
-        define_singleton_method m do |*args|
-          @connection.send(m, *args)
-        end
+      # default caching options
+      @cache_enabled = {
+        file: true,
+        command: false,
+      }
+
+      @cache_enabled.each_key do |type|
+        @cache[type] = {}
       end
     end
 
+    def clear_cache(type)
+      @cache[type.to_sym] = {}
+    end
+
     def file(path)
-      if @file_cache.key?(path)
-        @file_cache[path]
+      if @cache[:file].key?(path)
+        @cache[:file][path]
       else
-        @file_cache[path] = @connection.file(path)
+        @cache[:file][path] = @connection.file_via_connection(path)
       end
     end
 
     def run_command(cmd)
-      if @cmd_cache.key?(cmd)
-        @cmd_cache[cmd]
+      if @cache[:command].key?(cmd)
+        @cache[:command][cmd]
       else
-        @cmd_cache[cmd] = @connection.run_command(cmd)
+        @cache[:command][cmd] = @connection.run_command_via_connection(cmd)
       end
     end
   end

lib/train/plugins/transport.rb

@@ -39,13 +39,6 @@ def self.name(name)
       Train::Plugins.registry[name] = self
     end
 
-    # Create a cache connection to the target.
-    #
-    # @return [CacheConnection]
-    def cache_connection
-      CacheConnection.new(connection)
-    end
-
     private
 
     # @return [Logger] logger for reporting information

lib/train/transports/docker.rb

@@ -69,7 +69,7 @@ def close
       # nothing to do at the moment
     end
 
-    def file(path)
+    def file_via_connection(path)
       if os.aix?
         Train::File::Remote::Aix.new(self, path)
       elsif os.solaris?
@@ -79,7 +79,7 @@ def file(path)
       end
     end
 
-    def run_command(cmd)
+    def run_command_via_connection(cmd)
       cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?
       stdout, stderr, exit_status = @container.exec(
         [

lib/train/transports/local.rb

@@ -23,7 +23,7 @@ def initialize(options)
         @cmd_wrapper = CommandWrapper.load(self, options)
       end
 
-      def run_command(cmd)
+      def run_command_via_connection(cmd)
         cmd = @cmd_wrapper.run(cmd) unless @cmd_wrapper.nil?
         res = Mixlib::ShellOut.new(cmd)
         res.run_command
@@ -36,7 +36,7 @@ def local?
         true
       end
 
-      def file(path)
+      def file_via_connection(path)
         if os.windows?
           Train::File::Local::Windows.new(self, path)
         else

lib/train/transports/mock.rb

@@ -104,7 +104,7 @@ def command_not_found(cmd)
       mock_command(cmd, nil, nil, 1)
     end
 
-    def run_command(cmd)
+    def run_command_via_connection(cmd)
       @commands[cmd] ||
         @commands[Digest::SHA256.hexdigest cmd.to_s] ||
         command_not_found(cmd)
@@ -115,7 +115,7 @@ def file_not_found(path)
       File.new(self, path)
     end
 
-    def file(path)
+    def file_via_connection(path)
       @files[path] ||= file_not_found(path)
     end
 

lib/train/transports/ssh_connection.rb

@@ -55,7 +55,7 @@ def close
       @session = nil
     end
 
-    def file(path)
+    def file_via_connection(path)
       if os.aix?
         Train::File::Remote::Aix.new(self, path)
       elsif os.solaris?
@@ -68,7 +68,7 @@ def file(path)
     end
 
     # (see Base::Connection#run_command)
-    def run_command(cmd)
+    def run_command_via_connection(cmd)
       stdout = stderr = ''
       exit_status = nil
       cmd.dup.force_encoding('binary') if cmd.respond_to?(:force_encoding)

lib/train/transports/winrm_connection.rb

@@ -46,11 +46,11 @@ def close
       @session = nil
     end
 
-    def file(path)
+    def file_via_connection(path)
       Train::File::Remote::Windows.new(self, path)
     end
 
-    def run_command(command)
+    def run_command_via_connection(command)
       return if command.nil?
       logger.debug("[WinRM] #{self} (#{command})")
       out = ''

test/unit/plugins/cache_connection_test.rb

@@ -10,37 +10,52 @@
     it 'return new cache connection' do
       cache_connection.must_be_kind_of Train::Plugins::Transport::CacheConnection
       cache_connection.respond_to?(:run_command).must_equal true
+      cache_connection.respond_to?(:file).must_equal true
+    end
+
+    it 'default connection cache settings' do
+      cacher = connection.instance_variable_get(:@cacher)
+      cacher.cache_enabled[:file].must_equal true
+      cacher.cache_enabled[:command].must_equal false
     end
   end
 
-  describe 'check instance method fallback' do
-    it 'check os method' do
-      cache_connection.respond_to?(:os).must_equal true
-      cache_connection.respond_to?(:platform).must_equal true
+  describe 'disable/enable caching' do
+    it 'enable file cache' do
+      connection.enable_cache(:file)
+      cache_connection.cache_enabled[:file].must_equal true
     end
+  end
 
-    it 'check close method' do
-      cache_connection.respond_to?(:close).must_equal true
+  describe 'clear cache' do
+    it 'clear file cache' do
+      cache = cache_connection.instance_variable_get(:@cache)
+      cache[:file]['/tmp'] = 'test'
+      cache_connection.clear_cache(:file)
+      cache = cache_connection.instance_variable_get(:@cache)
+      cache[:file].must_equal({})
     end
   end
 
   describe 'load file' do
     it 'load file with caching' do
-      connection.expects(:file).once.returns('test_file')
+      connection.expects(:file_via_connection).once.returns('test_file')
       cache_connection.file('/tmp/test').must_equal('test_file')
       cache_connection.file('/tmp/test').must_equal('test_file')
       assert = { '/tmp/test' => 'test_file' }
-      cache_connection.instance_variable_get(:@file_cache).must_equal(assert)
+      cache = cache_connection.instance_variable_get(:@cache)
+      cache[:file].must_equal(assert)
     end
   end
 
   describe 'run command' do
     it 'run command with caching' do
-      connection.expects(:run_command).once.returns('test_user')
+      connection.expects(:run_command_via_connection).once.returns('test_user')
       cache_connection.run_command('whoami').must_equal('test_user')
       cache_connection.run_command('whoami').must_equal('test_user')
       assert = { 'whoami' => 'test_user' }
-      cache_connection.instance_variable_get(:@cmd_cache).must_equal(assert)
+      cache = cache_connection.instance_variable_get(:@cache)
+      cache[:command].must_equal(assert)
     end
   end
 end

test/unit/plugins/connection_test.rb

@@ -14,6 +14,10 @@
       proc { connection.run_command('') }.must_raise Train::ClientError
     end
 
+    it 'provides a run_command_via_connection method' do
+      proc { connection.run_command_via_connection('') }.must_raise Train::ClientError
+    end
+
     it 'provides an os method' do
       proc { connection.os }.must_raise Train::ClientError
     end
@@ -22,6 +26,10 @@
       proc { connection.file('') }.must_raise Train::ClientError
     end
 
+    it 'provides a file_via_connection method' do
+      proc { connection.file_via_connection('') }.must_raise Train::ClientError
+    end
+
     it 'provides a login command method' do
       proc { connection.login_command }.must_raise Train::ClientError
     end