Test ProxyFlightTime middleware

1 files changed, 31 insertions, 0 deletions

diff --git a/spec/lib/gitlab/middleware/proxy_flight_time_spec.rb b/spec/lib/gitlab/middleware/proxy_flight_time_spec.rb
new file mode 100644
index 00000000000..f59166d5c29
--- /dev/null
+++ b/spec/lib/gitlab/middleware/proxy_flight_time_spec.rb
@@ -0,0 +1,31 @@
+require 'spec_helper'
+
+describe Gitlab::Middleware::ProxyFlightTime do
+  let(:app) { double(:app) }
+  let(:middleware) { described_class.new(app) }
+  let(:env) { {} }
+  let(:transaction) { double(:transaction) }
+
+  before { expect(app).to receive(:call).with(env).and_return('yay') }
+
+  describe '#call' do
+    it 'calls the app when metrics are disabled' do
+      expect(Gitlab::Metrics).to receive(:current_transaction).and_return(nil)
+      expect(middleware.call(env)).to eq('yay')
+    end
+
+    context 'when metrics are enabled' do
+      before { allow(Gitlab::Metrics).to receive(:current_transaction).and_return(transaction) }
+
+      it 'calls the app when metrics are enabled but no timing header is found' do
+        expect(middleware.call(env)).to eq('yay')
+      end
+
+      it 'sets proxy_flight_time and calls the app when the header is present' do
+        env['HTTP_GITLAB_WORHORSE_PROXY_START'] = '123'
+        expect(transaction).to receive(:set).with(:proxy_flight_time, an_instance_of(Float))
+        expect(middleware.call(env)).to eq('yay')
+      end
+    end
+  end
+end