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
|
# 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
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) } }
let(:component_map) do
{
"application" => "test",
"endpoint_id" => "MarginaliaTestController#first_user",
"correlation_id" => correlation_id
}
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
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']
}
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' do
# We need to ensure that this runs through Sidekiq to take
# advantage of the middleware. There is a Rails bug that means we
# have to do some extra steps to make this happen:
# https://github.com/rails/rails/issues/37270#issuecomment-553927324
around do |example|
descendants = ActiveJob::Base.descendants + [ActiveJob::Base]
descendants.each(&:disable_test_adapter)
ActiveJob::Base.queue_adapter = :sidekiq
example.run
descendants.each { |a| a.queue_adapter = :test }
end
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
}
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)
request_env = Rack::MockRequest.env_for('/')
::Labkit::Correlation::CorrelationId.use_id(correlation_id) do
MarginaliaTestController.action(:first_user).call(request_env)
end
end
end
|