1 files changed, 133 insertions, 12 deletions

diff --git a/spec/lib/gitlab/workhorse_spec.rb b/spec/lib/gitlab/workhorse_spec.rb
index a1c2f7d667f..9bc1ebaebcb 100644
--- a/spec/lib/gitlab/workhorse_spec.rb
+++ b/spec/lib/gitlab/workhorse_spec.rb
@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Workhorse do
+RSpec.describe Gitlab::Workhorse, feature_category: :shared do
   let_it_be(:project) { create(:project, :repository) }
   let(:features) { { 'gitaly-feature-enforce-requests-limits' => 'true' } }
 
@@ -365,19 +365,72 @@ RSpec.describe Gitlab::Workhorse do
     end
   end
 
+  describe '.cleanup_key' do
+    let(:key) { 'test-key' }
+    let(:value) { 'test-value' }
+
+    subject(:cleanup_key) { described_class.cleanup_key(key) }
+
+    shared_examples 'cleans up key' do |redis = Gitlab::Redis::Workhorse|
+      before do
+        described_class.set_key_and_notify(key, value)
+      end
+
+      it 'deletes the key' do
+        expect { cleanup_key }
+          .to change { redis.with { |c| c.exists?(key) } }.from(true).to(false)
+      end
+    end
+
+    it_behaves_like 'cleans up key'
+
+    context 'when workhorse migration feature flags are disabled' do
+      before do
+        stub_feature_flags(
+          use_primary_and_secondary_stores_for_workhorse: false,
+          use_primary_store_as_default_for_workhorse: false
+        )
+      end
+
+      it_behaves_like 'cleans up key', Gitlab::Redis::SharedState
+    end
+
+    context 'when either workhorse migration feature flags are enabled' do
+      context 'when use_primary_and_secondary_stores_for_workhorse is enabled' do
+        before do
+          stub_feature_flags(
+            use_primary_store_as_default_for_workhorse: false
+          )
+        end
+
+        it_behaves_like 'cleans up key'
+      end
+
+      context 'when use_primary_store_as_default_for_workhorse is enabled' do
+        before do
+          stub_feature_flags(
+            use_primary_and_secondary_stores_for_workhorse: false
+          )
+        end
+
+        it_behaves_like 'cleans up key'
+      end
+    end
+  end
+
   describe '.set_key_and_notify' do
     let(:key) { 'test-key' }
     let(:value) { 'test-value' }
 
     subject { described_class.set_key_and_notify(key, value, overwrite: overwrite) }
 
-    shared_examples 'set and notify' do
+    shared_examples 'set and notify' do |redis = Gitlab::Redis::Workhorse|
       it 'set and return the same value' do
         is_expected.to eq(value)
       end
 
       it 'set and notify' do
-        expect(Gitlab::Redis::SharedState).to receive(:with).and_call_original
+        expect(redis).to receive(:with).and_call_original
         expect_any_instance_of(::Redis).to receive(:publish)
           .with(described_class::NOTIFICATION_PREFIX + 'test-key', "test-value")
 
@@ -389,6 +442,39 @@ RSpec.describe Gitlab::Workhorse do
       let(:overwrite) { true }
 
       it_behaves_like 'set and notify'
+
+      context 'when workhorse migration feature flags are disabled' do
+        before do
+          stub_feature_flags(
+            use_primary_and_secondary_stores_for_workhorse: false,
+            use_primary_store_as_default_for_workhorse: false
+          )
+        end
+
+        it_behaves_like 'set and notify', Gitlab::Redis::SharedState
+      end
+
+      context 'when either workhorse migration feature flags are enabled' do
+        context 'when use_primary_and_secondary_stores_for_workhorse is enabled' do
+          before do
+            stub_feature_flags(
+              use_primary_store_as_default_for_workhorse: false
+            )
+          end
+
+          it_behaves_like 'set and notify'
+        end
+
+        context 'when use_primary_store_as_default_for_workhorse is enabled' do
+          before do
+            stub_feature_flags(
+              use_primary_and_secondary_stores_for_workhorse: false
+            )
+          end
+
+          it_behaves_like 'set and notify'
+        end
+      end
     end
 
     context 'when we set an existing key' do
@@ -519,18 +605,53 @@ RSpec.describe Gitlab::Workhorse do
   describe '.send_dependency' do
     let(:headers) { { Accept: 'foo', Authorization: 'Bearer asdf1234' } }
     let(:url) { 'https://foo.bar.com/baz' }
+    let(:upload_method) { nil }
+    let(:upload_url) { nil }
+    let(:upload_headers) { {} }
+    let(:upload_config) { { method: upload_method, headers: upload_headers, url: upload_url }.compact_blank! }
 
-    subject { described_class.send_dependency(headers, url) }
+    subject { described_class.send_dependency(headers, url, upload_config: upload_config) }
 
-    it 'sets the header correctly', :aggregate_failures do
-      key, command, params = decode_workhorse_header(subject)
+    shared_examples 'setting the header correctly' do |ensure_upload_config_field: nil|
+      it 'sets the header correctly' do
+        key, command, params = decode_workhorse_header(subject)
+        expected_params = {
+          'Headers' => headers.transform_values { |v| Array.wrap(v) },
+          'Url' => url,
+          'UploadConfig' => {
+            'Method' => upload_method,
+            'Url' => upload_url,
+            'Headers' => upload_headers.transform_values { |v| Array.wrap(v) }
+          }.compact_blank!
+        }
+        expected_params.compact_blank!
 
-      expect(key).to eq("Gitlab-Workhorse-Send-Data")
-      expect(command).to eq("send-dependency")
-      expect(params).to eq({
-        'Header' => headers,
-        'Url' => url
-      }.deep_stringify_keys)
+        expect(key).to eq("Gitlab-Workhorse-Send-Data")
+        expect(command).to eq("send-dependency")
+        expect(params).to eq(expected_params.deep_stringify_keys)
+
+        expect(params.dig('UploadConfig', ensure_upload_config_field)).to be_present if ensure_upload_config_field
+      end
+    end
+
+    it_behaves_like 'setting the header correctly'
+
+    context 'overriding the method' do
+      let(:upload_method) { 'PUT' }
+
+      it_behaves_like 'setting the header correctly', ensure_upload_config_field: 'Method'
+    end
+
+    context 'overriding the upload url' do
+      let(:upload_url) { 'https://test.dev' }
+
+      it_behaves_like 'setting the header correctly', ensure_upload_config_field: 'Url'
+    end
+
+    context 'with upload headers set' do
+      let(:upload_headers) { { 'Private-Token' => '1234567890' } }
+
+      it_behaves_like 'setting the header correctly', ensure_upload_config_field: 'Headers'
     end
   end