Slave lag throttler symbolize config keys

lib/lhm/throttler/slave_lag.rb

@@ -91,9 +91,10 @@ class Slave
 
       attr_reader :host, :connection
 
-      def initialize(host, get_config=nil)
+      def initialize(host, connection_config = nil)
         @host = host
-        @connection = client(config(get_config))
+        @connection_config = prepare_connection_config(connection_config)
+        @connection = client(@connection_config)
       end
 
       def slave_hosts
@@ -108,17 +109,26 @@ def lag
 
       def client(config)
         begin
-          Lhm.logger.info "Connecting to #{@host} on database: #{config['database']}"
+          Lhm.logger.info "Connecting to #{@host} on database: #{config[:database]}"
           Mysql2::Client.new(config)
         rescue Mysql2::Error => e
           Lhm.logger.info "Error connecting to #{@host}: #{e}"
           nil
         end
       end
 
-      def config(get_config)
-        config = get_config ? get_config.call : ActiveRecord::Base.connection_pool.spec.config.dup
-        config['host'] = @host
+      def prepare_connection_config(config_proc)
+        config = if config_proc
+          if config_proc.respond_to?(:call) # if we get a proc
+            config_proc.call
+          else
+            raise ArgumentError, "Expected #{config_proc.inspect} to respond to `call`"
+          end
+        else # otherwise default to ActiveRecord provided config
+          ActiveRecord::Base.connection_pool.spec.config.dup
+        end
+        config.deep_symbolize_keys!
+        config[:host] = @host
         config
       end
 

spec/unit/throttler/slave_lag_spec.rb

@@ -36,9 +36,10 @@
   include UnitHelper
 
   before :each do
-    def get_config
-      lambda { {'username' => 'user', 'password' => 'pw', 'database' => 'db'} }
-    end
+    @logs = StringIO.new
+    Lhm.logger = Logger.new(@logs)
+
+    @dummy_mysql_client_config = lambda { {'username' => 'user', 'password' => 'pw', 'database' => 'db'} }
   end
 
   describe "#client" do
@@ -52,24 +53,36 @@ def initialize(config)
 
     describe 'on connection error' do
       it 'logs and returns nil' do
-        test_client = lambda { |config|
-          TestMysql2Client.new(config)
-        }
-        Mysql2::Client.stubs(:new).returns(test_client) do
-          assert_send([Lhm.logger, :info, "Error connecting to slave: connection error"])
-          assert_nil(Lhm::Throttler::Slave.new('slave', lambda { {} }).connection)
-        end
+        assert_nil(Lhm::Throttler::Slave.new('slave', @dummy_mysql_client_config).connection)
+
+        log_messages = @logs.string.lines
+        assert_equal(2, log_messages.length)
+        assert log_messages[0].include? "Connecting to slave on database: db"
+        assert log_messages[1].include? "Error connecting to slave: Unknown MySQL server host 'slave'"
       end
     end
 
     describe 'with proper config' do
       it "creates a new Mysql2::Client" do
-        client_assertion = lambda { |config|
-          assert_equal(config, {'username' => 'user', 'password' => 'pw', 'database' => 'db', 'host' => 'slave'})
-        }
-        Mysql2::Client.stubs(:new).returns(client_assertion) do
-          Lhm::Throttler::Slave.new('slave', get_config)
-        end
+        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
+      end
+    end
+
+    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)))
+
+        Mysql2::Client.stubs(:new).returns(mock())
+
+        assert Lhm::Throttler::Slave.new('slave').connection
+
+        log_messages = @logs.string.lines
+        assert_equal(1, log_messages.length)
+        assert log_messages[0].include? "Connecting to slave on database: db"
       end
     end
   end
@@ -86,7 +99,7 @@ def self.query(query)
         end
       end
 
-      @slave = Lhm::Throttler::Slave.new('slave', get_config)
+      @slave = Lhm::Throttler::Slave.new('slave', @dummy_mysql_client_config)
       @slave.instance_variable_set(:@connection, Connection)
 
       class StoppedConnection
@@ -95,7 +108,7 @@ def self.query(query)
         end
       end
 
-      @stopped_slave = Lhm::Throttler::Slave.new('stopped_slave', get_config)
+      @stopped_slave = Lhm::Throttler::Slave.new('stopped_slave', @dummy_mysql_client_config)
       @stopped_slave.instance_variable_set(:@connection, StoppedConnection)
     end
 
@@ -124,7 +137,7 @@ def self.query(query)
         Lhm::Throttler::Slave.any_instance.stubs(:client).returns(client)
         Lhm::Throttler::Slave.any_instance.stubs(:config).returns([])
 
-        slave = Lhm::Throttler::Slave.new('slave', get_config)
+        slave = Lhm::Throttler::Slave.new('slave', @dummy_mysql_client_config)
         assert_send([Lhm.logger, :info, "Unable to connect and/or query slave: error"])
         assert_equal(0, slave.lag)
       end
@@ -210,7 +223,7 @@ def @throttler.max_current_slave_lag
         client.stubs(:query).raises(Mysql2::Error, "Can't connect to MySQL server")
         Lhm::Throttler::Slave.any_instance.stubs(:client).returns(client)
 
-        Lhm::Throttler::Slave.any_instance.stubs(:config).returns([])
+        Lhm::Throttler::Slave.any_instance.stubs(:prepare_connection_config).returns([])
         Lhm::Throttler::Slave.any_instance.stubs(:slave_hosts).returns(['1.1.1.2'])
         @throttler.stubs(:master_slave_hosts).returns(['1.1.1.1'])
 
@@ -237,7 +250,7 @@ def @throttler.master_slave_hosts
         class TestSlave
           attr_reader :host, :connection
 
-          def initialize(host, get_config)
+          def initialize(host, _)
             @host = host
             @connection = 'conn' if @host
           end