Fix flaky GCS test

Shopify · Nov 21, 2024 · b46eb0f · b46eb0f
1 parent 96ea13e
commit b46eb0f
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 12 deletions.
diff --git a/lib/app_profiler/storage/google_cloud_storage.rb b/lib/app_profiler/storage/google_cloud_storage.rb
@@ -50,12 +50,6 @@ def enqueue_upload(profile)
           end
         end
 
-        def reset_queue # for testing
-          init_queue
-          @process_queue_thread&.kill
-          @process_queue_thread = nil
-        end
-
         private
 
         def mutex

diff --git a/test/app_profiler/middleware_test.rb b/test/app_profiler/middleware_test.rb
@@ -352,7 +352,7 @@ class MiddlewareTest < TestCase
         assert(response[1]["X-Profile-Async"])
       end
     ensure
-      AppProfiler::Storage::GoogleCloudStorage.reset_queue # kill the background thread and reset the queue
+      reset_process_queue_thread # kill the background thread and reset the queue
       AppProfiler.storage = old_storage
     end
 

diff --git a/test/app_profiler/storage/google_cloud_storage_test.rb b/test/app_profiler/storage/google_cloud_storage_test.rb
@@ -9,9 +9,7 @@ class GoogleCloudStorageTest < AppProfiler::TestCase
       TEST_BUCKET_NAME = "app-profile-test"
       TEST_FILE_URL = "https://www.example.com/uploaded.json"
 
-      def teardown
-        GoogleCloudStorage.reset_queue
-      end
+      teardown(:reset_process_queue_thread)
 
       test "upload file" do
         profile = profile_from_stackprof
@@ -114,9 +112,8 @@ def teardown
       end
 
       test "process_queue_thread is alive after first upload" do
-        th = GoogleCloudStorage.instance_variable_get(:@process_queue_thread)
+        reset_process_queue_thread
 
-        refute(th&.alive?)
         GoogleCloudStorage.enqueue_upload(profile_from_stackprof)
         th = GoogleCloudStorage.instance_variable_get(:@process_queue_thread)
         assert(th.alive?)

diff --git a/test/test_helper.rb b/test/test_helper.rb
@@ -94,5 +94,11 @@ def with_yarn_setup(app)
     ensure
       app.instance_variable_set(:@yarn_initialized, old_yarn_setup)
     end
+
+    def reset_process_queue_thread
+      Storage::GoogleCloudStorage.send(:init_queue)
+      Storage::GoogleCloudStorage.instance_variable_get(:@process_queue_thread)&.kill
+      Storage::GoogleCloudStorage.instance_variable_set(:@process_queue_thread, nil)
+    end
   end
 end