blob: 01646bc76efbbe4e3815cd65fa94d3c93c7c2a3e (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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
|
