Use token from redis in gcp project billing worker

2 files changed, 46 insertions, 15 deletions

diff --git a/app/workers/check_gcp_project_billing_worker.rb b/app/workers/check_gcp_project_billing_worker.rb
index 42aa6b39d86..784d17c4654 100644
--- a/app/workers/check_gcp_project_billing_worker.rb
+++ b/app/workers/check_gcp_project_billing_worker.rb
@@ -1,14 +1,23 @@
+require 'securerandom'
+
 class CheckGcpProjectBillingWorker
   include ApplicationWorker
   include ClusterQueue
 
   LEASE_TIMEOUT = 15.seconds.to_i
 
+  def self.generate_redis_token_key
+    SecureRandom.uuid
+  end
+
   def self.redis_shared_state_key_for(token)
     "gitlab:gcp:#{token.hash}:billing_enabled"
   end
 
-  def perform(token)
+  def perform(token_key)
+    return unless token_key
+
+    token = get_token(token_key)
     return unless token
     return unless try_obtain_lease_for(token)
 
@@ -20,6 +29,10 @@ class CheckGcpProjectBillingWorker
 
   private
 
+  def get_token(token_key)
+    Gitlab::Redis::SharedState.with { |redis| redis.get(token_key) }
+  end
+
   def try_obtain_lease_for(token)
     Gitlab::ExclusiveLease
       .new("check_gcp_project_billing_worker:#{token.hash}", timeout: LEASE_TIMEOUT)
diff --git a/spec/workers/check_gcp_project_billing_worker_spec.rb b/spec/workers/check_gcp_project_billing_worker_spec.rb
index cdb114749ee..b1687cc95c6 100644
--- a/spec/workers/check_gcp_project_billing_worker_spec.rb
+++ b/spec/workers/check_gcp_project_billing_worker_spec.rb
@@ -3,33 +3,51 @@ require 'spec_helper'
 describe CheckGcpProjectBillingWorker do
   describe '.perform' do
     let(:token) { 'bogustoken' }
-    subject { described_class.new.perform(token) }
+    subject { described_class.new.perform('token_key') }
 
-    context 'when there is no lease' do
+    context 'when there is a token in redis' do
       before do
-        allow_any_instance_of(described_class).to receive(:try_obtain_lease_for).and_return('randomuuid')
+        allow_any_instance_of(described_class).to receive(:get_token).and_return(token)
       end
 
-      it 'calls the service' do
-        expect(CheckGcpProjectBillingService).to receive_message_chain(:new, :execute).and_return([double])
+      context 'when there is no lease' do
+        before do
+          allow_any_instance_of(described_class).to receive(:try_obtain_lease_for).and_return('randomuuid')
+        end
 
-        subject
+        it 'calls the service' do
+          expect(CheckGcpProjectBillingService).to receive_message_chain(:new, :execute).and_return([double])
+
+          subject
+        end
+
+        it 'stores billing status in redis' do
+          redis_double = double
+
+          expect(CheckGcpProjectBillingService).to receive_message_chain(:new, :execute).and_return([double])
+          expect(Gitlab::Redis::SharedState).to receive(:with).and_yield(redis_double)
+          expect(redis_double).to receive(:set).with(described_class.redis_shared_state_key_for(token), anything)
+
+          subject
+        end
       end
 
-      it 'stores billing status in redis' do
-        redis_double = double
+      context 'when there is a lease' do
+        before do
+          allow_any_instance_of(described_class).to receive(:try_obtain_lease_for).and_return(false)
+        end
 
-        expect(CheckGcpProjectBillingService).to receive_message_chain(:new, :execute).and_return([double])
-        expect(Gitlab::Redis::SharedState).to receive(:with).and_yield(redis_double)
-        expect(redis_double).to receive(:set).with(described_class.redis_shared_state_key_for(token), anything)
+        it 'does not call the service' do
+          expect(CheckGcpProjectBillingService).not_to receive(:new)
 
-        subject
+          subject
+        end
       end
     end
 
-    context 'when there is a lease' do
+    context 'when there is no token in redis' do
       before do
-        allow_any_instance_of(described_class).to receive(:try_obtain_lease_for).and_return(false)
+        allow_any_instance_of(described_class).to receive(:get_token).and_return(nil)
       end
 
       it 'does not call the service' do