lib/switchman/shackles.rb in switchman-1.4.3 vs lib/switchman/shackles.rb in switchman-1.4.4

@@ -1,10 +1,10 @@
 module Switchman
   module Shackles
     module ClassMethods
-      def ensure_handler
-        raise "This should not be called with switchman installed"
+      def self.prepended(klass)
+        klass.send(:remove_method, :ensure_handler)
       end
 
       # drops the save_handler and ensure_handler calls from the vanilla
       # Shackles' implementation.
       def activate!(environment)
@@ -21,15 +21,8 @@
         old_environment = activate!(environment)
         yield
       ensure
         activate!(old_environment)
       end
-    end
-
-    def self.included(klass)
-      klass.extend(ClassMethods)
-      klass.singleton_class.send(:remove_method, :ensure_handler)
-      klass.singleton_class.send(:remove_method, :activate!)
-      klass.singleton_class.send(:remove_method, :activate)
     end
   end
 end