1 files changed, 69 insertions, 0 deletions
diff --git a/spec/lib/api/ci/helpers/runner_spec.rb b/spec/lib/api/ci/helpers/runner_spec.rb
new file mode 100644
index 00000000000..99f2db544a5
--- /dev/null
+++ b/spec/lib/api/ci/helpers/runner_spec.rb
@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe API::Ci::Helpers::Runner do
+  let(:helper) { Class.new { include API::Ci::Helpers::Runner }.new }
+
+  before do
+    allow(helper).to receive(:env).and_return({})
+  end
+
+  describe '#current_job' do
+    let(:build) { create(:ci_build, :running) }
+
+    it 'handles sticking of a build when a build ID is specified' do
+      allow(helper).to receive(:params).and_return(id: build.id)
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .to receive(:stick_or_unstick)
+        .with({}, :build, build.id)
+
+      helper.current_job
+    end
+
+    it 'does not handle sticking if no build ID was specified' do
+      allow(helper).to receive(:params).and_return({})
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .not_to receive(:stick_or_unstick)
+
+      helper.current_job
+    end
+
+    it 'returns the build if one could be found' do
+      allow(helper).to receive(:params).and_return(id: build.id)
+
+      expect(helper.current_job).to eq(build)
+    end
+  end
+
+  describe '#current_runner' do
+    let(:runner) { create(:ci_runner, token: 'foo') }
+
+    it 'handles sticking of a runner if a token is specified' do
+      allow(helper).to receive(:params).and_return(token: runner.token)
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .to receive(:stick_or_unstick)
+        .with({}, :runner, runner.token)
+
+      helper.current_runner
+    end
+
+    it 'does not handle sticking if no token was specified' do
+      allow(helper).to receive(:params).and_return({})
+
+      expect(Gitlab::Database::LoadBalancing::RackMiddleware)
+        .not_to receive(:stick_or_unstick)
+
+      helper.current_runner
+    end
+
+    it 'returns the runner if one could be found' do
+      allow(helper).to receive(:params).and_return(token: runner.token)
+
+      expect(helper.current_runner).to eq(runner)
+    end
+  end
+end