diff --git a/lib/migration_tools/migration_extension.rb b/lib/migration_tools/migration_extension.rb
index e41ac76..f1079c6 100644
--- a/lib/migration_tools/migration_extension.rb
+++ b/lib/migration_tools/migration_extension.rb
@@ -1,14 +1,15 @@
 module MigrationTools
   module MigrationExtension
+    module ClassMethods
+      attr_accessor :migration_group
 
-    attr_accessor :migration_group
+      def group(arg = nil)
+        unless MigrationTools::MIGRATION_GROUPS.member?(arg.to_s)
+          raise "Invalid group \"#{arg.to_s}\" - valid groups are #{MigrationTools::MIGRATION_GROUPS.inspect}"
+        end
 
-    def group(arg = nil)
-      unless MigrationTools::MIGRATION_GROUPS.member?(arg.to_s)
-        raise "Invalid group \"#{arg.to_s}\" - valid groups are #{MigrationTools::MIGRATION_GROUPS.inspect}"
+        self.migration_group = arg.to_s
       end
-
-      self.migration_group = arg.to_s
     end
 
     def migrate_with_forced_groups(direction)
@@ -21,14 +22,15 @@ def migrate_with_forced_groups(direction)
 end
 
 ActiveRecord::Migration.class_eval do
-  extend MigrationTools::MigrationExtension
-  class << self
-    alias_method :migrate_without_forced_groups, :migrate
-    alias_method :migrate, :migrate_with_forced_groups
-  end
+  extend MigrationTools::MigrationExtension::ClassMethods
+  include MigrationTools::MigrationExtension
+
+  alias_method :migrate_without_forced_groups, :migrate
+  alias_method :migrate, :migrate_with_forced_groups
 
   def migration_group
     self.class.migration_group
   end
 end
+
 ActiveRecord::MigrationProxy.delegate :migration_group, :to => :migration
diff --git a/test/test_migration_tools.rb b/test/test_migration_tools.rb
index 1e8bd98..4b7af8b 100644
--- a/test/test_migration_tools.rb
+++ b/test/test_migration_tools.rb
@@ -16,6 +16,10 @@
     @task = MigrationTools::Tasks.new
   end
 
+  after do
+    MigrationTools.instance_variable_set('@forced', false)
+  end
+
   def migrations
     [ Alpha, Beta, Delta, Kappa ]
   end
@@ -56,17 +60,15 @@ def migration_proxy(m)
 
   it "forcing" do
     assert !MigrationTools.forced?
-    Kappa.migrate("up")
-
     MigrationTools.forced!
     assert MigrationTools.forced?
 
-    Alpha.migrate("up")
+    @task.migrator(0).run
 
     begin
-      Kappa.migrate("up")
+      @task.migrator(3).run
       fail "You should not be able to run migrations without groups in forced mode"
-    rescue RuntimeError => e
+    rescue => e
       assert e.message =~ /Cowardly refusing/
     end
   end