2 files changed, 26 insertions, 2 deletions
diff --git a/spec/uploaders/object_storage_spec.rb b/spec/uploaders/object_storage_spec.rb
index fb51ae4c02c..4503288e410 100644
--- a/spec/uploaders/object_storage_spec.rb
+++ b/spec/uploaders/object_storage_spec.rb
@@ -732,4 +732,26 @@ describe ObjectStorage do
       end
     end
   end
+
+  describe '#retrieve_from_store!' do
+    [:group, :project, :user].each do |model|
+      context "for #{model}s" do
+        let(:models) { create_list(model, 3, :with_avatar).map(&:reload) }
+        let(:avatars) { models.map(&:avatar) }
+
+        it 'batches fetching uploads from the database' do
+          # Ensure that these are all created and fully loaded before we start
+          # running queries for avatars
+          models
+
+          expect { avatars }.not_to exceed_query_limit(1)
+        end
+
+        it 'fetches a unique upload for each model' do
+          expect(avatars.map(&:url).uniq).to eq(avatars.map(&:url))
+          expect(avatars.map(&:upload).uniq).to eq(avatars.map(&:upload))
+        end
+      end
+    end
+  end
 end
diff --git a/spec/uploaders/workers/object_storage/background_move_worker_spec.rb b/spec/uploaders/workers/object_storage/background_move_worker_spec.rb
index b34f427fd8a..95813d15e52 100644
--- a/spec/uploaders/workers/object_storage/background_move_worker_spec.rb
+++ b/spec/uploaders/workers/object_storage/background_move_worker_spec.rb
@@ -125,8 +125,10 @@ describe ObjectStorage::BackgroundMoveWorker do
 
           it "migrates file to remote storage" do
             perform
+            project.reload
+            BatchLoader::Executor.clear_current
 
-            expect(project.reload.avatar.file_storage?).to be_falsey
+            expect(project.avatar).not_to be_file_storage
           end
         end
 
@@ -137,7 +139,7 @@ describe ObjectStorage::BackgroundMoveWorker do
           it "migrates file to remote storage" do
             perform
 
-            expect(project.reload.avatar.file_storage?).to be_falsey
+            expect(project.reload.avatar).not_to be_file_storage
           end
         end
       end