3 files changed, 65 insertions, 17 deletions
diff --git a/activerecord/test/cases/identity_map_test.rb b/activerecord/test/cases/identity_map_test.rb
index 959c303d88..649715fbb5 100644
--- a/activerecord/test/cases/identity_map_test.rb
+++ b/activerecord/test/cases/identity_map_test.rb
@@ -1,5 +1,4 @@
 require "cases/helper"
-require "active_support/log_subscriber/test_helper"
 
 require 'models/developer'
 require 'models/project'
@@ -28,8 +27,6 @@ class IdentityMapTest < ActiveRecord::TestCase
     :developers_projects, :computers, :authors, :author_addresses,
     :posts, :tags, :taggings, :comments, :subscribers
 
-  include ActiveSupport::LogSubscriber::TestHelper
-
   ##############################################################################
   # Basic tests checking if IM is functioning properly on basic find operations#
   ##############################################################################
@@ -386,20 +383,6 @@ class IdentityMapTest < ActiveRecord::TestCase
     assert_not_nil post.title
   end
 
-  def test_log
-    # FIXME: Can't seem to figure out why it isn't logging in test, works fine in a real app
-    pending "LogSubscriber wonkiness"
-    @old_logger = ActiveRecord::Base.logger
-    ActiveRecord::LogSubscriber.attach_to(:active_record)
-
-    Post.find 1
-    Post.find 1
-    assert_match(/From Identity Map/, @logger.logged(:debug).last)
-  ensure
-    ActiveRecord::LogSubscriber.log_subscribers.pop
-    ActiveRecord::Base.logger = @old_logger
-  end
-
 # Currently AR is not allowing changing primary key (see Persistence#update)
 # So we ignore it. If this changes, this test needs to be uncommented.
 #  def test_updating_of_pkey
diff --git a/activerecord/test/cases/log_subscriber_test.rb b/activerecord/test/cases/log_subscriber_test.rb
index 5f55299065..c6c6079490 100644
--- a/activerecord/test/cases/log_subscriber_test.rb
+++ b/activerecord/test/cases/log_subscriber_test.rb
@@ -1,11 +1,14 @@
 require "cases/helper"
 require "models/developer"
+require "models/post"
 require "active_support/log_subscriber/test_helper"
 
 class LogSubscriberTest < ActiveRecord::TestCase
   include ActiveSupport::LogSubscriber::TestHelper
   include ActiveSupport::BufferedLogger::Severity
 
+  fixtures :posts
+
   def setup
     @old_logger = ActiveRecord::Base.logger
     @using_identity_map = ActiveRecord::IdentityMap.enabled?
@@ -91,4 +94,13 @@ class LogSubscriberTest < ActiveRecord::TestCase
   def test_initializes_runtime
     Thread.new { assert_equal 0, ActiveRecord::LogSubscriber.runtime }.join
   end
+
+  def test_log
+    ActiveRecord::IdentityMap.use do
+      Post.find 1
+      Post.find 1
+    end
+    wait
+    assert_match(/From Identity Map/, @logger.logged(:debug).last)
+  end
 end
diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb
index 287f7e255b..b2e40c6b22 100644
--- a/activerecord/test/cases/query_cache_test.rb
+++ b/activerecord/test/cases/query_cache_test.rb
@@ -10,6 +10,59 @@ class QueryCacheTest < ActiveRecord::TestCase
 
   def setup
     Task.connection.clear_query_cache
+    ActiveRecord::Base.connection.disable_query_cache!
+  end
+
+  def test_middleware_delegates
+    called = false
+    mw = ActiveRecord::QueryCache.new lambda { |env|
+      called = true
+    }
+    mw.call({})
+    assert called, 'middleware should delegate'
+  end
+
+  def test_middleware_caches
+    mw = ActiveRecord::QueryCache.new lambda { |env|
+      Task.find 1
+      Task.find 1
+      assert_equal 1, ActiveRecord::Base.connection.query_cache.length
+    }
+    mw.call({})
+  end
+
+  def test_cache_enabled_during_call
+    assert !ActiveRecord::Base.connection.query_cache_enabled, 'cache off'
+
+    mw = ActiveRecord::QueryCache.new lambda { |env|
+      assert ActiveRecord::Base.connection.query_cache_enabled, 'cache on'
+    }
+    mw.call({})
+  end
+
+  def test_cache_on_during_body_write
+    streaming = Class.new do
+      def each
+        yield ActiveRecord::Base.connection.query_cache_enabled
+      end
+    end
+
+    mw = ActiveRecord::QueryCache.new lambda { |env|
+      [200, {}, streaming.new]
+    }
+    body = mw.call({}).last
+    body.each { |x| assert x, 'cache should be on' }
+    body.close
+    assert !ActiveRecord::Base.connection.query_cache_enabled, 'cache disabled'
+  end
+
+  def test_cache_off_after_close
+    mw = ActiveRecord::QueryCache.new lambda { |env| }
+    body = mw.call({}).last
+
+    assert ActiveRecord::Base.connection.query_cache_enabled, 'cache enabled'
+    body.close
+    assert !ActiveRecord::Base.connection.query_cache_enabled, 'cache disabled'
   end
 
   def test_find_queries