lhm-shopify 3.5.0 → 3.5.4

data/spec/test_helper.rb

@@ -10,6 +10,8 @@ require 'minitest/autorun'
 require 'minitest/spec'
 require 'minitest/mock'
 require 'mocha/minitest'
+require 'after_do'
+require 'byebug'
 require 'pathname'
 require 'lhm'
 
@@ -17,6 +19,8 @@ $project = Pathname.new(File.dirname(__FILE__) + '/..').cleanpath
 $spec = $project.join('spec')
 $fixtures = $spec.join('fixtures')
 
+$db_name = 'test'
+
 require 'active_record'
 require 'mysql2'
 
@@ -43,3 +47,20 @@ end
 def throttler
   Lhm::Throttler::Time.new(:stride => 100)
 end
+
+def init_test_db
+  db_config = YAML.load_file(File.expand_path(File.dirname(__FILE__)) + '/integration/database.yml')
+  conn = Mysql2::Client.new(
+    :host => '127.0.0.1',
+    :username => db_config['master']['user'],
+    :password => db_config['master']['password'],
+    :port => db_config['master']['port']
+  )
+
+  conn.query("DROP DATABASE IF EXISTS #{$db_name}")
+  conn.query("CREATE DATABASE #{$db_name}")
+end
+
+init_test_db
+
+

data/spec/unit/chunk_insert_spec.rb

@@ -4,17 +4,22 @@
 require File.expand_path(File.dirname(__FILE__)) + '/unit_helper'
 
 require 'lhm/chunk_insert'
+require 'lhm/connection'
 
 describe Lhm::ChunkInsert do
   before(:each) do
-    @connection = stub(:connection)
+    ar_connection = mock()
+    ar_connection.stubs(:execute).returns([["dummy"]])
+    @connection = Lhm::Connection.new(connection: ar_connection, options: {reconnect_with_consistent_host: false})
     @origin = Lhm::Table.new('foo')
     @destination = Lhm::Table.new('bar')
   end
 
   describe "#sql" do
     describe "when migration has no conditions" do
-      before { @migration = Lhm::Migration.new(@origin, @destination) }
+      before do
+        @migration = Lhm::Migration.new(@origin, @destination)
+      end
 
       it "uses a simple where clause" do
         assert_equal(

data/spec/unit/chunker_spec.rb

@@ -12,11 +12,15 @@ require 'lhm/connection'
 describe Lhm::Chunker do
   include UnitHelper
 
+  EXPECTED_RETRY_FLAGS_CHUNKER = {:should_retry => true, :log_prefix => "Chunker"}
+  EXPECTED_RETRY_FLAGS_CHUNK_INSERT = {:should_retry => true, :log_prefix => "ChunkInsert"}
+
   before(:each) do
     @origin = Lhm::Table.new('foo')
     @destination = Lhm::Table.new('bar')
     @migration = Lhm::Migration.new(@origin, @destination)
     @connection = mock()
+    @connection.stubs(:execute).returns([["dummy"]])
     # This is a poor man's stub
     @throttler = Object.new
     def @throttler.run
@@ -38,11 +42,11 @@ describe Lhm::Chunker do
         5
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 4/),{}).returns(7)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 4/),{}).returns(21)
-      @connection.expects(:update).with(regexp_matches(/between 1 and 7/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 8 and 10/),{}).returns(2)
-      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),{}).returns([])
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 4/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(7)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 4/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(21)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 7/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 8 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       @chunker.run
     end
@@ -53,17 +57,17 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),{}).returns(2)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 1/),{}).returns(4)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 5 order by id limit 1 offset 1/),{}).returns(6)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 7 order by id limit 1 offset 1/),{}).returns(8)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 1/),{}).returns(10)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(2)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(4)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 5 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(6)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 7 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(8)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(10)
 
-      @connection.expects(:update).with(regexp_matches(/between 1 and 2/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 3 and 4/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 5 and 6/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 7 and 8/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),{}).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 2/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 3 and 4/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 5 and 6/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 7 and 8/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
 
       @chunker.run
     end
@@ -80,17 +84,17 @@ describe Lhm::Chunker do
         end
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),{}).returns(2)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 2/),{}).returns(5)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 2/),{}).returns(8)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 2/),{}).returns(nil)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(2)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 3 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(5)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(8)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 9 order by id limit 1 offset 2/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(nil)
 
-      @connection.expects(:update).with(regexp_matches(/between 1 and 2/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 3 and 5/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 6 and 8/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),{}).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 2/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 3 and 5/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 6 and 8/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 9 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
 
-      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),{}).returns([])
+      @connection.expects(:execute).twice.with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       @chunker.run
     end
@@ -100,8 +104,8 @@ describe Lhm::Chunker do
                                                            :start     => 1,
                                                            :limit     => 1)
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 0/),{}).returns(nil)
-      @connection.expects(:update).with(regexp_matches(/between 1 and 1/),{}).returns(1)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 0/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(nil)
+      @connection.expects(:update).with(regexp_matches(/between 1 and 1/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
 
       @chunker.run
     end
@@ -114,17 +118,17 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 2 order by id limit 1 offset 1/),{}).returns(3)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 4 order by id limit 1 offset 1/),{}).returns(5)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 1/),{}).returns(7)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 1/),{}).returns(9)
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 10 order by id limit 1 offset 1/),{}).returns(nil)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 2 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(3)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 4 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(5)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 6 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(7)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 8 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(9)
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 10 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(nil)
 
-      @connection.expects(:update).with(regexp_matches(/between 2 and 3/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 4 and 5/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 6 and 7/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 8 and 9/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/between 10 and 10/),{}).returns(1)
+      @connection.expects(:update).with(regexp_matches(/between 2 and 3/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 4 and 5/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 6 and 7/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 8 and 9/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(2)
+      @connection.expects(:update).with(regexp_matches(/between 10 and 10/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
 
       @chunker.run
     end
@@ -138,9 +142,9 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/where \(foo.created_at > '2013-07-10' or foo.baz = 'quux'\) and `foo`/),{}).returns(1)
-      @connection.expects(:execute).with(regexp_matches(/show warnings/),{}).returns([])
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(2)
+      @connection.expects(:update).with(regexp_matches(/where \(foo.created_at > '2013-07-10' or foo.baz = 'quux'\) and `foo`/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
+      @connection.expects(:execute).with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       def @migration.conditions
         "where foo.created_at > '2013-07-10' or foo.baz = 'quux'"
@@ -158,9 +162,9 @@ describe Lhm::Chunker do
         2
       end
 
-      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),{}).returns(2)
-      @connection.expects(:update).with(regexp_matches(/inner join bar on foo.id = bar.foo_id and/),{}).returns(1)
-      @connection.expects(:execute).with(regexp_matches(/show warnings/),{}).returns([])
+      @connection.expects(:select_value).with(regexp_matches(/where id >= 1 order by id limit 1 offset 1/),EXPECTED_RETRY_FLAGS_CHUNKER).returns(2)
+      @connection.expects(:update).with(regexp_matches(/inner join bar on foo.id = bar.foo_id and/),EXPECTED_RETRY_FLAGS_CHUNK_INSERT).returns(1)
+      @connection.expects(:execute).with(regexp_matches(/show warnings/),EXPECTED_RETRY_FLAGS_CHUNKER).returns([])
 
       def @migration.conditions
         'inner join bar on foo.id = bar.foo_id'

data/spec/unit/connection_spec.rb

@@ -1,4 +1,5 @@
 require 'lhm/connection'
+require 'lhm/proxysql_helper'
 
 describe Lhm::Connection do
 
@@ -12,10 +13,15 @@ describe Lhm::Connection do
   it "Should find use calling file as prefix" do
     ar_connection = mock()
     ar_connection.stubs(:execute).raises(LOCK_WAIT).then.returns(true)
+    ar_connection.stubs(:active?).returns(true)
 
-    connection = Lhm::Connection.new(connection: ar_connection)
+    connection = Lhm::Connection.new(connection: ar_connection, options: {
+      retriable: {
+        base_interval: 0
+      }
+    })
 
-    connection.execute("SHOW TABLES", { base_interval: 0 })
+    connection.execute("SHOW TABLES", should_retry: true)
 
     log_messages = @logs.string.split("\n")
     assert_equal(1, log_messages.length)
@@ -27,10 +33,16 @@ describe Lhm::Connection do
     ar_connection.stubs(:execute).raises(LOCK_WAIT)
                  .then.raises(LOCK_WAIT)
                  .then.returns(true)
+    ar_connection.stubs(:active?).returns(true)
 
-    connection = Lhm::Connection.new(connection: ar_connection)
+    connection = Lhm::Connection.new(connection: ar_connection, options: {
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
 
-    connection.execute("SHOW TABLES", { base_interval: 0, tries: 3 })
+    connection.execute("SHOW TABLES", should_retry: true)
 
     log_messages = @logs.string.split("\n")
     assert_equal(2, log_messages.length)
@@ -41,10 +53,16 @@ describe Lhm::Connection do
     ar_connection.stubs(:update).raises(LOCK_WAIT)
                  .then.raises(LOCK_WAIT)
                  .then.returns(1)
+    ar_connection.stubs(:active?).returns(true)
 
-    connection = Lhm::Connection.new(connection: ar_connection)
+    connection = Lhm::Connection.new(connection: ar_connection, options: {
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
 
-    val = connection.update("SHOW TABLES", { base_interval: 0, tries: 3 })
+    val = connection.update("SHOW TABLES", should_retry: true)
 
     log_messages = @logs.string.split("\n")
     assert_equal val, 1
@@ -56,13 +74,38 @@ describe Lhm::Connection do
     ar_connection.stubs(:select_value).raises(LOCK_WAIT)
                  .then.raises(LOCK_WAIT)
                  .then.returns("dummy")
+    ar_connection.stubs(:active?).returns(true)
 
-    connection = Lhm::Connection.new(connection: ar_connection)
+    connection = Lhm::Connection.new(connection: ar_connection, options: {
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
 
-    val = connection.select_value("SHOW TABLES", { base_interval: 0, tries: 3 })
+    val = connection.select_value("SHOW TABLES", should_retry: true)
 
     log_messages = @logs.string.split("\n")
     assert_equal val, "dummy"
     assert_equal(2, log_messages.length)
   end
+
+  it "Queries should be tagged with ProxySQL tag if reconnect_with_consistent_host is enabled" do
+    ar_connection = mock()
+    ar_connection.expects(:public_send).with(:select_value, "SHOW TABLES #{Lhm::ProxySQLHelper::ANNOTATION}").returns("dummy")
+    ar_connection.stubs(:execute).times(4).returns([["dummy"]])
+    ar_connection.stubs(:active?).returns(true)
+
+    connection = Lhm::Connection.new(connection: ar_connection, options: {
+      reconnect_with_consistent_host: true,
+      retriable: {
+        base_interval: 0,
+        tries: 3
+      }
+    })
+
+    val = connection.select_value("SHOW TABLES", should_retry: true)
+
+    assert_equal val, "dummy"
+  end
 end

data/spec/unit/entangler_spec.rb

@@ -63,42 +63,94 @@ describe Lhm::Entangler do
     it 'should retry trigger creation when it hits a lock wait timeout' do
       tries = 1
       ar_connection = mock()
-      ar_connection.expects(:execute).times(tries).raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
-
-      connection = Lhm::Connection.new(connection: ar_connection)
-
-      @entangler = Lhm::Entangler.new(@migration, connection, retriable: {base_interval: 0, tries: tries})
+      ar_connection.stubs(:execute)
+                       .returns([["dummy"]], [["dummy"]], [["dummy"]])
+                       .then
+                       .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+      ar_connection.stubs(:active?).returns(true)
+
+      connection = Lhm::Connection.new(connection: ar_connection, options: {
+        reconnect_with_consistent_host: true,
+        retriable: {
+          base_interval: 0,
+          tries: tries
+        }
+      })
+
+      @entangler = Lhm::Entangler.new(@migration, connection)
 
       assert_raises(Mysql2::Error) { @entangler.before }
     end
 
     it 'should not retry trigger creation with other mysql errors' do
       ar_connection = mock()
-      ar_connection.expects(:execute).once.raises(Mysql2::Error, 'The MySQL server is running with the --read-only option so it cannot execute this statement.')
-      connection = Lhm::Connection.new(connection: ar_connection)
-
-      @entangler = Lhm::Entangler.new(@migration, connection, retriable: {base_interval: 0})
+      ar_connection.stubs(:execute)
+                   .returns([["dummy"]], [["dummy"]], [["dummy"]])
+                   .then
+                   .raises(Mysql2::Error, 'The MySQL server is running with the --read-only option so it cannot execute this statement.')
+      ar_connection.stubs(:active?).returns(true)
+      connection = Lhm::Connection.new(connection: ar_connection, options: {
+        reconnect_with_consistent_host: true,
+        retriable: {
+          base_interval: 0
+        },
+      })
+
+      @entangler = Lhm::Entangler.new(@migration, connection)
       assert_raises(Mysql2::Error) { @entangler.before }
     end
 
     it 'should succesfully finish after retrying' do
       ar_connection = mock()
-      ar_connection.stubs(:execute).raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction').then.returns(true)
-
-      connection = Lhm::Connection.new(connection: ar_connection)
-
-      @entangler = Lhm::Entangler.new(@migration, connection, retriable: {base_interval: 0})
+      ar_connection.stubs(:execute)
+                   .returns([["dummy"]], [["dummy"]], [["dummy"]])
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]])
+      ar_connection.stubs(:active?).returns(true)
+
+      connection = Lhm::Connection.new(connection: ar_connection, options: {
+        reconnect_with_consistent_host: true,
+        retriable: {
+          base_interval: 0
+        },
+      })
+
+      @entangler = Lhm::Entangler.new(@migration, connection)
 
       assert @entangler.before
     end
 
     it 'should retry as many times as specified by configuration' do
       ar_connection = mock()
-      ar_connection.expects(:execute).times(5).raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
-
-      connection = Lhm::Connection.new(connection: ar_connection)
-
-      @entangler = Lhm::Entangler.new(@migration, connection, retriable: {tries: 5, base_interval: 0})
+      ar_connection.stubs(:execute)
+                   .returns([["dummy"]], [["dummy"]], [["dummy"]]) # initial
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]]) # reconnect 1
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]])  # reconnect 2
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')
+                   .then
+                   .returns([["dummy"]])  # reconnect 3
+                   .then
+                   .raises(Mysql2::Error, 'Lock wait timeout exceeded; try restarting transaction')  # final error
+      ar_connection.stubs(:active?).returns(true)
+
+      connection = Lhm::Connection.new(connection: ar_connection, options: {
+        reconnect_with_consistent_host: true,
+        retriable: {
+          tries: 5,
+          base_interval: 0
+        },
+      })
+
+      @entangler = Lhm::Entangler.new(@migration, connection)
 
       assert_raises(Mysql2::Error) { @entangler.before }
     end

data/spec/unit/lhm_spec.rb

@@ -26,4 +26,21 @@ describe Lhm do
       value(Lhm.logger.instance_eval { @logdev }.dev.path).must_equal 'omg.ponies'
     end
   end
+
+  describe 'api' do
+
+    before(:each) do
+      @connection = mock()
+    end
+
+    it 'should create a new connection when calling setup' do
+      Lhm.setup(@connection)
+      value(Lhm.connection).must_be_kind_of(Lhm::Connection)
+    end
+
+    it 'should create a new connection when none is created' do
+      ActiveRecord::Base.stubs(:connection).returns(@connection)
+      value(Lhm.connection).must_be_kind_of(Lhm::Connection)
+    end
+  end
 end

data/spec/unit/throttler/slave_lag_spec.rb

@@ -39,7 +39,7 @@ describe Lhm::Throttler::Slave do
     @logs = StringIO.new
     Lhm.logger = Logger.new(@logs)
 
-    @dummy_mysql_client_config = lambda { {'username' => 'user', 'password' => 'pw', 'database' => 'db'} }
+    @dummy_mysql_client_config = lambda { { 'username' => 'user', 'password' => 'pw', 'database' => 'db' } }
   end
 
   describe "#client" do
@@ -64,7 +64,7 @@ describe Lhm::Throttler::Slave do
 
     describe 'with proper config' do
       it "creates a new Mysql2::Client" do
-        expected_config = {username: 'user', password: 'pw', database: 'db', host: 'slave'}
+        expected_config = { username: 'user', password: 'pw', database: 'db', host: 'slave' }
         Mysql2::Client.stubs(:new).with(expected_config).returns(mock())
 
         assert Lhm::Throttler::Slave.new('slave', @dummy_mysql_client_config).connection
@@ -73,8 +73,12 @@ describe Lhm::Throttler::Slave do
 
     describe 'with active record config' do
       it 'logs and creates client' do
-        active_record_config = {username: 'user', password: 'pw', database: 'db'}
-        ActiveRecord::Base.stubs(:connection_pool).returns(stub(spec: stub(config: active_record_config)))
+        active_record_config = { username: 'user', password: 'pw', database: 'db' }
+        if ActiveRecord::VERSION::MAJOR > 6 || ActiveRecord::VERSION::MAJOR == 6 && ActiveRecord::VERSION::MINOR >= 1
+          ActiveRecord::Base.stubs(:connection_pool).returns(stub(db_config: stub(configuration_hash: active_record_config)))
+        else
+          ActiveRecord::Base.stubs(:connection_pool).returns(stub(spec: stub(config: active_record_config)))
+        end
 
         Mysql2::Client.stubs(:new).returns(mock())
 
@@ -92,9 +96,9 @@ describe Lhm::Throttler::Slave do
       class Connection
         def self.query(query)
           if query == Lhm::Throttler::Slave::SQL_SELECT_MAX_SLAVE_LAG
-            [{'Seconds_Behind_Master' => 20}]
+            [{ 'Seconds_Behind_Master' => 20 }]
           elsif query == Lhm::Throttler::Slave::SQL_SELECT_SLAVE_HOSTS
-            [{'host' => '1.1.1.1:80'}]
+            [{ 'host' => '1.1.1.1:80' }]
           end
         end
       end
@@ -104,7 +108,7 @@ describe Lhm::Throttler::Slave do
 
       class StoppedConnection
         def self.query(query)
-          [{'Seconds_Behind_Master' => nil}]
+          [{ 'Seconds_Behind_Master' => nil }]
         end
       end
 
@@ -138,7 +142,7 @@ describe Lhm::Throttler::Slave do
         Lhm::Throttler::Slave.any_instance.stubs(:config).returns([])
 
         slave = Lhm::Throttler::Slave.new('slave', @dummy_mysql_client_config)
-        assert_send([Lhm.logger, :info, "Unable to connect and/or query slave: error"])
+        Logger.any_instance.expects(:info).with("Unable to connect and/or query slave: Can't connect to MySQL server")
         assert_equal(0, slave.lag)
       end
     end
@@ -286,7 +290,7 @@ describe Lhm::Throttler::SlaveLag do
       describe 'with the :check_only option' do
         describe 'with a callable argument' do
           before do
-            check_only = lambda {{'host' => '1.1.1.3'}}
+            check_only = lambda { { 'host' => '1.1.1.3' } }
             @throttler = Lhm::Throttler::SlaveLag.new :check_only => check_only
           end
 
@@ -300,6 +304,7 @@ describe Lhm::Throttler::SlaveLag do
         describe 'with a non-callable argument' do
           before do
             @throttler = Lhm::Throttler::SlaveLag.new :check_only => 'I cannot be called'
+
             def @throttler.master_slave_hosts
               ['1.1.1.1', '1.1.1.4']
             end