Rakefile in elasticsearch-api-0.4.3 vs Rakefile in elasticsearch-api-0.4.4

@@ -1,7 +1,9 @@
 require "bundler/gem_tasks"
 
+__current__ = Pathname( File.expand_path('..', __FILE__) )
+
 desc "Run unit tests"
 task :default => 'test:unit'
 task :test    => 'test:unit'
 
 # ----- Test tasks ------------------------------------------------------------
@@ -25,22 +27,65 @@
     test.test_files = FileList["test/unit/**/*_test.rb"]
     # test.verbose = true
     # test.warning = true
   end
 
-  Rake::TestTask.new(:integration) do |test|
-    Rake::Task['test:ci_reporter'].invoke if ENV['CI']
-    test.libs << 'lib' << 'test'
-    test.test_files = FileList["test/integration/yaml_test_runner.rb", "test/integration/**/*_test.rb"]
+  desc "Run integration tests"
+  task :integration do
+    require 'elasticsearch'
+
+    current_branch = `git --git-dir=#{__current__.join('../support/elasticsearch/.git')} --work-tree=#{__current__.join('../support/elasticsearch')} branch --no-color`.split("\n").select { |b| b =~ /^\*/ }.first.gsub(/^\*\s*/, '')
+
+    # Define the task
+    t = Rake::TestTask.new(:integration) do |test|
+      Rake::Task['test:ci_reporter'].invoke if ENV['CI']
+      test.libs << 'lib' << 'test'
+      test.test_files = FileList["test/integration/yaml_test_runner.rb", "test/integration/**/*_test.rb"]
+    end
+
+    begin
+      client = Elasticsearch::Client.new host: "localhost:#{ENV['TEST_CLUSTER_PORT'] || 9250}"
+      es_version_info = client.info['version']
+      build_hash      = es_version_info['build_hash']
+      cluster_running = true
+    rescue Faraday::Error::ConnectionFailed => e
+      STDERR.puts "[!] Test cluster not running?"
+      cluster_running = false
+    end
+
+    checkout_server_build = ENV['TEST_NO_SPEC_CHECKOUT'].nil? || true
+    checkout_build_hash   = ENV['TEST_BUILD_HASH'] || build_hash
+
+    begin
+      # Checkout the branch corresponding to the build
+      sh "git --git-dir=#{__current__.join('../support/elasticsearch/.git')} --work-tree=#{__current__.join('../support/elasticsearch')} checkout #{checkout_build_hash} --force --quiet" if cluster_running && checkout_server_build
+
+      # Path to the API specs
+      ENV['TEST_REST_API_SPEC'] = __current__.join('../support/elasticsearch/rest-api-spec/test/').to_s
+
+      # Run the task
+      args = [t.ruby_opts_string, t.run_code, t.file_list_string, t.option_list].join(' ')
+
+      ruby args do |ok, status|
+        if !ok && status.respond_to?(:signaled?) && status.signaled?
+          raise SignalException.new(status.termsig)
+        elsif !ok
+          fail "Command failed with status (#{status.exitstatus}): " + "[ruby #{args}]"
+        end
+      end
+
+    ensure
+      sh "git --git-dir=#{__current__.join('../support/elasticsearch/.git')} --work-tree=#{__current__.join('../support/elasticsearch')} checkout #{current_branch}", verbose: false if cluster_running && checkout_server_build
+    end
   end
 
   Rake::TestTask.new(:all) do |test|
     Rake::Task['test:ci_reporter'].invoke if ENV['CI']
     test.libs << 'lib' << 'test'
     test.test_files = FileList["test/unit/**/*_test.rb", "test/integration/**/*_test.rb", "test/integration/yaml_test_runner.rb"]
   end
 
-  namespace :server do
+  namespace :cluster do
     desc "Start Elasticsearch nodes for tests"
     task :start do
       $LOAD_PATH << File.expand_path('../../elasticsearch-transport/lib', __FILE__) << File.expand_path('../test', __FILE__)
       require 'elasticsearch/extensions/test/cluster'
       Elasticsearch::Extensions::Test::Cluster.start