From 748ee58cbb2d79516c540e3b4b5ca67eb72648c6 Mon Sep 17 00:00:00 2001
From: Gustavo Caso <gustavo.caso@datadoghq.com>
Date: Tue, 21 Mar 2023 14:21:13 +0100
Subject: [PATCH] Apply feedback

---
 lib/datadog/core/remote/worker.rb       |  5 ++-
 sig/datadog/core/remote/worker.rbs      |  4 +--
 spec/datadog/core/remote/worker_spec.rb | 42 ++++++++++++++-----------
 spec/spec_helper.rb                     |  3 +-
 4 files changed, 30 insertions(+), 24 deletions(-)

diff --git a/lib/datadog/core/remote/worker.rb b/lib/datadog/core/remote/worker.rb
index 9916ca94e5e..6e2d702ef11 100644
--- a/lib/datadog/core/remote/worker.rb
+++ b/lib/datadog/core/remote/worker.rb
@@ -47,7 +47,10 @@ def stop
 
           thread = @thr
 
-          thread.kill if thread
+          if thread
+            thread.kill
+            thread.join
+          end
 
           @started = false
           @stopping = false
diff --git a/sig/datadog/core/remote/worker.rbs b/sig/datadog/core/remote/worker.rbs
index f9ebde5a56b..6a7561812fd 100644
--- a/sig/datadog/core/remote/worker.rbs
+++ b/sig/datadog/core/remote/worker.rbs
@@ -8,9 +8,9 @@ module Datadog
         attr_reader thr: Thread?
         attr_reader interval: Integer
         attr_reader mutex: ::Thread::Mutex
-        attr_reader block: (^() -> untyped)
+        attr_reader block: (^() -> void)
 
-        def initialize: (interval: Integer) { () -> untyped } -> void
+        def initialize: (interval: Integer) { () -> void } -> void
 
         def start: () -> void
 
diff --git a/spec/datadog/core/remote/worker_spec.rb b/spec/datadog/core/remote/worker_spec.rb
index 341c4121e6e..0e71135ee40 100644
--- a/spec/datadog/core/remote/worker_spec.rb
+++ b/spec/datadog/core/remote/worker_spec.rb
@@ -4,6 +4,9 @@
 require 'datadog/core/remote/worker'
 
 RSpec.describe Datadog::Core::Remote::Worker do
+  let(:task) { proc { 1 + 1 } }
+  subject(:worker) { described_class.new(interval: 1, &task) }
+
   describe '#initialize' do
     it 'raises ArgumentError when no block is provided' do
       expect do
@@ -12,37 +15,38 @@
     end
   end
 
-  subject(:worker) do
-    described_class.new(interval: 1) do
-      1 + 1
-    end
-  end
-
   describe '#start' do
+    after { worker.stop }
+
     it 'mark worker as started' do
       expect(worker).not_to be_started
       worker.start
       expect(worker).to be_started
-      worker.stop
     end
 
     it 'acquire and release lock' do
-      expect(worker).to receive(:acquire_lock)
-      expect(worker).to receive(:release_lock)
+      expect(worker).to receive(:acquire_lock).at_least(:once)
+      expect(worker).to receive(:release_lock).at_least(:once)
       worker.start
     end
 
-    it 'execute block when started' do
-      result = []
-      queue = Queue.new
-      queue_worker = described_class.new(interval: 1) do
-        result << queue.pop
+    context 'execute block when started' do
+      let(:result) { [] }
+      let(:queue) { Queue.new }
+      let(:task) do
+        proc do
+          value = 1
+          result << value
+          queue << value
+        end
+      end
+
+      it 'runs block' do
+        worker.start
+        # Wait for the work task to execute once
+        queue.pop
+        expect(result).to eq([1])
       end
-      queue_worker.start
-      # Unblock worker thread
-      queue << 1
-      expect(result).to eq([1])
-      queue_worker.stop
     end
   end
 
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index fb618d2bc18..e57d0b9d3f3 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -168,8 +168,7 @@
         # suite output unreadable.
         if example.file_path.start_with?(
           './spec/datadog/core/workers/',
-          './spec/ddtrace/workers/',
-          './spec/datadog/core/remote/worker_spec.rb'
+          './spec/ddtrace/workers/'
         )
           puts # Add newline so we get better output when the progress formatter is being used
           RSpec.warning("FIXME: #{example.file_path}:#{example.metadata[:line_number]} is leaking threads")