Add latest changes from gitlab-org/gitlab@master

spec/lib/gitlab/sidekiq_middleware/worker_context/client_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::SidekiqMiddleware::WorkerContext::Client do
+  let(:worker_class) do
+    Class.new do
+      def self.name
+        'TestWithContextWorker'
+      end
+
+      include ApplicationWorker
+
+      def self.job_for_args(args)
+        jobs.find { |job| job['args'] == args }
+      end
+
+      def perform(*args)
+      end
+    end
+  end
+
+  before do
+    stub_const('TestWithContextWorker', worker_class)
+  end
+
+  describe "#call" do
+    it 'applies a context for jobs scheduled in batch' do
+      user_per_job = { 'job1' => build_stubbed(:user, username: 'user-1'),
+                       'job2' => build_stubbed(:user, username: 'user-2') }
+
+      TestWithContextWorker.bulk_perform_async_with_contexts(
+        %w(job1 job2),
+        arguments_proc: -> (name) { [name, 1, 2, 3] },
+        context_proc: -> (name) { { user: user_per_job[name] } }
+      )
+
+      job1 = TestWithContextWorker.job_for_args(['job1', 1, 2, 3])
+      job2 = TestWithContextWorker.job_for_args(['job2', 1, 2, 3])
+
+      expect(job1['meta.user']).to eq(user_per_job['job1'].username)
+      expect(job2['meta.user']).to eq(user_per_job['job2'].username)
+    end
+  end
+end

spec/lib/gitlab/sidekiq_middleware_spec.rb

@@ -110,6 +110,14 @@ describe Gitlab::SidekiqMiddleware do
     let(:queue) { 'default' }
     let(:redis_pool) { Sidekiq.redis_pool }
     let(:middleware_expected_args) { [worker_class_arg, job, queue, redis_pool] }
+    let(:expected_middlewares) do
+      [
+        Gitlab::SidekiqStatus::ClientMiddleware,
+        Gitlab::SidekiqMiddleware::ClientMetrics,
+        Gitlab::SidekiqMiddleware::WorkerContext::Client,
+        Labkit::Middleware::Sidekiq::Client
+      ]
+    end
     before do
       described_class.client_configurator.call(chain)
@@ -120,8 +128,9 @@ describe Gitlab::SidekiqMiddleware do
       # this will prevent the full middleware chain from being executed.
       # This test ensures that this does not happen
       it "invokes the chain" do
-        expect_any_instance_of(Gitlab::SidekiqStatus::ClientMiddleware).to receive(:call).with(*middleware_expected_args).once.and_call_original
-        expect_any_instance_of(Labkit::Middleware::Sidekiq::Client).to receive(:call).with(*middleware_expected_args).once.and_call_original
+        expected_middlewares do |middleware|
+          expect_any_instance_of(middleware).to receive(:call).with(*middleware_expected_args).once.ordered.and_call_original
+        end
         expect { |b| chain.invoke(worker_class_arg, job, queue, redis_pool, &b) }.to yield_control.once
       end

spec/support/db_cleaner.rb

 # frozen_string_literal: true
 module DbCleaner
-  def delete_from_all_tables!(except: nil)
+  def delete_from_all_tables!(except: [])
+    except << 'ar_internal_metadata'
+
     DatabaseCleaner.clean_with(:deletion, cache_tables: false, except: except)
   end

spec/workers/concerns/worker_context_spec.rb

@@ -5,7 +5,11 @@ require 'spec_helper'
 describe WorkerContext do
   let(:worker) do
     Class.new do
-      include WorkerContext
+      def self.name
+        "TestWorker"
+      end
+
+      include ApplicationWorker
     end
   end
@@ -24,6 +28,78 @@ describe WorkerContext do
     end
   end
+  shared_examples 'tracking bulk scheduling contexts' do
+    describe "context contents" do
+      before do
+        # stub clearing the contexts, so we can check what's inside
+        allow(worker).to receive(:batch_context=).and_call_original
+        allow(worker).to receive(:batch_context=).with(nil)
+      end
+
+      it 'keeps track of the context per key to schedule' do
+        subject
+
+        expect(worker.context_for_arguments(["hello"])).to be_a(Gitlab::ApplicationContext)
+      end
+
+      it 'does not share contexts across threads' do
+        t1_context = nil
+        t2_context = nil
+
+        Thread.new do
+          subject
+
+          t1_context = worker.context_for_arguments(["hello"])
+        end.join
+        Thread.new do
+          t2_context = worker.context_for_arguments(["hello"])
+        end.join
+
+        expect(t1_context).to be_a(Gitlab::ApplicationContext)
+        expect(t2_context).to be_nil
+      end
+    end
+
+    it 'clears the contexts' do
+      subject
+
+      expect(worker.__send__(:batch_context)).to be_nil
+    end
+  end
+
+  describe '.bulk_perform_async_with_contexts' do
+    subject do
+      worker.bulk_perform_async_with_contexts(%w(hello world),
+                                              context_proc: -> (_) { { user: build_stubbed(:user) } },
+                                              arguments_proc: -> (word) { word })
+    end
+
+    it 'calls bulk_perform_async with the arguments' do
+      expect(worker).to receive(:bulk_perform_async).with([["hello"], ["world"]])
+
+      subject
+    end
+
+    it_behaves_like 'tracking bulk scheduling contexts'
+  end
+
+  describe '.bulk_perform_in_with_contexts' do
+    subject do
+      worker.bulk_perform_in_with_contexts(10.minutes,
+                                           %w(hello world),
+                                           context_proc: -> (_) { { user: build_stubbed(:user) } },
+                                           arguments_proc: -> (word) { word })
+    end
+
+    it 'calls bulk_perform_in with the arguments and delay' do
+      expect(worker).to receive(:bulk_perform_in).with(10.minutes, [["hello"], ["world"]])
+
+      subject
+    end
+
+    it_behaves_like 'tracking bulk scheduling contexts'
+  end
+
   describe '#with_context' do
     it 'allows modifying context when the job is running' do
       worker.new.with_context(user: build_stubbed(:user, username: 'jane-doe')) do