lib/airbrussh/rake/context.rb in airbrussh-1.4.2 vs lib/airbrussh/rake/context.rb in airbrussh-1.5.0

@@ -11,10 +11,14 @@
     # task in the middle of an SSHKit thread, which causes Context to get very
     # confused. It such scenarios Context is not reliable and may return `nil`
     # for the `position` of a command.
     #
     class Context
+      class << self
+        attr_accessor :current_task_name
+      end
+
       def initialize(config=Airbrussh.configuration)
         @history = []
         @enabled = config.monkey_patch_rake
         self.class.install_monkey_patch if enabled?
       end
@@ -45,13 +49,25 @@
       # task. May be `nil` in certain multi-threaded scenarios, so be careful!
       def position(command)
         history.index(command.to_s)
       end
 
-      class << self
-        attr_accessor :current_task_name
+      if Object.respond_to?(:prepend)
+        module Patch
+          def execute(args=nil)
+            ::Airbrussh::Rake::Context.current_task_name = name.to_s
+            super
+          end
+        end
 
-        def install_monkey_patch
+        def self.install_monkey_patch
+          require "rake"
+          return if ::Rake::Task.include?(::Airbrussh::Rake::Context::Patch)
+
+          ::Rake::Task.prepend(::Airbrussh::Rake::Context::Patch)
+        end
+      else
+        def self.install_monkey_patch
           require "rake"
           return if ::Rake::Task.instance_methods.include?(:_airbrussh_execute)
 
           ::Rake::Task.class_exec do
             alias_method :_airbrussh_execute, :execute