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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
# frozen_string_literal: true
require 'spec_helper'
RSpec.describe 'Marginalia spec' do
class MarginaliaTestController < ApplicationController
skip_before_action :authenticate_user!, :check_two_factor_requirement
def first_user
User.first
render body: nil
end
def first_ci_pipeline
Ci::Pipeline.first
render body: nil
end
private
[:auth_user, :current_user, :set_experimentation_subject_id_cookie, :signed_in?].each do |method|
define_method(method) {}
end
end
class MarginaliaTestJob
include Sidekiq::Worker
def perform
Gitlab::ApplicationContext.with_context(caller_id: self.class.name) do
User.first
end
end
end
class MarginaliaTestMailer < ApplicationMailer
def first_user
User.first
end
end
describe 'For rails web requests' do
let(:correlation_id) { SecureRandom.uuid }
let(:recorded) { ActiveRecord::QueryRecorder.new { make_request(correlation_id, :first_user) } }
let(:component_map) do
{
"application" => "test",
"endpoint_id" => "MarginaliaTestController#first_user",
"correlation_id" => correlation_id,
"db_config_name" => "main"
}
end
it 'generates a query that includes the component and value' do
component_map.each do |component, value|
expect(recorded.log.last).to include("#{component}:#{value}")
end
end
context 'when using CI database' do
let(:recorded) { ActiveRecord::QueryRecorder.new { make_request(correlation_id, :first_ci_pipeline) } }
let(:component_map) do
{
"application" => "test",
"endpoint_id" => "MarginaliaTestController#first_ci_pipeline",
"correlation_id" => correlation_id,
"db_config_name" => 'ci'
}
end
before do
skip_if_multiple_databases_not_setup(:ci)
end
it 'generates a query that includes the component and value' do
component_map.each do |component, value|
expect(recorded.log.last).to include("#{component}:#{value}")
end
end
end
end
describe 'for Sidekiq worker jobs' do
around do |example|
with_sidekiq_server_middleware do |chain|
chain.add Labkit::Middleware::Sidekiq::Context::Server
chain.add Marginalia::SidekiqInstrumentation::Middleware
Marginalia.application_name = "sidekiq"
example.run
end
end
after(:all) do
MarginaliaTestJob.clear
end
before do
MarginaliaTestJob.perform_async
end
let(:sidekiq_job) { MarginaliaTestJob.jobs.first }
let(:recorded) { ActiveRecord::QueryRecorder.new { MarginaliaTestJob.drain } }
let(:component_map) do
{
"application" => "sidekiq",
"endpoint_id" => "MarginaliaTestJob",
"correlation_id" => sidekiq_job['correlation_id'],
"jid" => sidekiq_job['jid'],
"db_config_name" => "main"
}
end
it 'generates a query that includes the component and value' do
component_map.each do |component, value|
expect(recorded.log.last).to include("#{component}:#{value}")
end
end
describe 'for ActionMailer delivery jobs', :sidekiq_mailers do
let(:delivery_job) { MarginaliaTestMailer.first_user.deliver_later }
let(:recorded) do
ActiveRecord::QueryRecorder.new do
Sidekiq::Worker.drain_all
end
end
let(:component_map) do
{
"application" => "sidekiq",
"endpoint_id" => "ActionMailer::MailDeliveryJob",
"jid" => delivery_job.job_id,
"db_config_name" => "main"
}
end
it 'generates a query that includes the component and value' do
component_map.each do |component, value|
expect(recorded.log.last).to include("#{component}:#{value}")
end
end
end
end
def make_request(correlation_id, action_name)
request_env = Rack::MockRequest.env_for('/')
::Labkit::Correlation::CorrelationId.use_id(correlation_id) do
MarginaliaTestController.action(action_name).call(request_env)
end
end
end
|
