diff options
Diffstat (limited to 'spec/lib/gitlab/database/load_balancing/setup_spec.rb')
-rw-r--r-- | spec/lib/gitlab/database/load_balancing/setup_spec.rb | 119 |
1 files changed, 119 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/load_balancing/setup_spec.rb b/spec/lib/gitlab/database/load_balancing/setup_spec.rb new file mode 100644 index 00000000000..01646bc76ef --- /dev/null +++ b/spec/lib/gitlab/database/load_balancing/setup_spec.rb @@ -0,0 +1,119 @@ +# frozen_string_literal: true + +require 'spec_helper' + +RSpec.describe Gitlab::Database::LoadBalancing::Setup do + describe '#setup' do + it 'sets up the load balancer' do + setup = described_class.new(ActiveRecord::Base) + + expect(setup).to receive(:disable_prepared_statements) + expect(setup).to receive(:setup_load_balancer) + expect(setup).to receive(:setup_service_discovery) + + setup.setup + end + end + + describe '#disable_prepared_statements' do + it 'disables prepared statements and reconnects to the database' do + config = double( + :config, + configuration_hash: { host: 'localhost' }, + env_name: 'test', + name: 'main' + ) + model = double(:model, connection_db_config: config) + + expect(ActiveRecord::DatabaseConfigurations::HashConfig) + .to receive(:new) + .with('test', 'main', { host: 'localhost', prepared_statements: false }) + .and_call_original + + # HashConfig doesn't implement its own #==, so we can't directly compare + # the expected value with a pre-defined one. + expect(model) + .to receive(:establish_connection) + .with(an_instance_of(ActiveRecord::DatabaseConfigurations::HashConfig)) + + described_class.new(model).disable_prepared_statements + end + end + + describe '#setup_load_balancer' do + it 'sets up the load balancer' do + model = Class.new(ActiveRecord::Base) + setup = described_class.new(model) + config = Gitlab::Database::LoadBalancing::Configuration.new(model) + lb = instance_spy(Gitlab::Database::LoadBalancing::LoadBalancer) + + allow(lb).to receive(:configuration).and_return(config) + + expect(Gitlab::Database::LoadBalancing::LoadBalancer) + .to receive(:new) + .with(setup.configuration) + .and_return(lb) + + setup.setup_load_balancer + + expect(model.connection.load_balancer).to eq(lb) + expect(model.sticking) + .to be_an_instance_of(Gitlab::Database::LoadBalancing::Sticking) + end + end + + describe '#setup_service_discovery' do + context 'when service discovery is disabled' do + it 'does nothing' do + expect(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .not_to receive(:new) + + described_class.new(ActiveRecord::Base).setup_service_discovery + end + end + + context 'when service discovery is enabled' do + it 'immediately performs service discovery' do + model = ActiveRecord::Base + setup = described_class.new(model) + sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery) + lb = model.connection.load_balancer + + allow(setup.configuration) + .to receive(:service_discovery_enabled?) + .and_return(true) + + allow(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .to receive(:new) + .with(lb, setup.configuration.service_discovery) + .and_return(sv) + + expect(sv).to receive(:perform_service_discovery) + expect(sv).not_to receive(:start) + + setup.setup_service_discovery + end + + it 'starts service discovery if needed' do + model = ActiveRecord::Base + setup = described_class.new(model, start_service_discovery: true) + sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery) + lb = model.connection.load_balancer + + allow(setup.configuration) + .to receive(:service_discovery_enabled?) + .and_return(true) + + allow(Gitlab::Database::LoadBalancing::ServiceDiscovery) + .to receive(:new) + .with(lb, setup.configuration.service_discovery) + .and_return(sv) + + expect(sv).to receive(:perform_service_discovery) + expect(sv).to receive(:start) + + setup.setup_service_discovery + end + end + end +end |