lib/draper/test/test_unit_integration.rb in draper-1.0.0.beta2 vs lib/draper/test/test_unit_integration.rb in draper-1.0.0.beta3

@@ -1,9 +1,18 @@
 require "rake/testtask"
-require "rails/test_unit/sub_test_task"
 
+klass = nil
+
+if Rails.version[0,3] == "3.0"
+  require 'rails/test_unit/railtie'
+  klass = Rake::TestTask
+else 
+  require "rails/test_unit/sub_test_task"
+  klass = Rails::SubTestTask
+end
+
 namespace :test do
-  Rails::SubTestTask.new(:decorators => "test:prepare") do |t|
+  klass.new(:decorators => "test:prepare") do |t|
     t.libs << "test"
     t.pattern = "test/decorators/**/*_test.rb"
   end
 end