Memoize Aws::SNS::Clinet.new

cookpad · Jul 11, 2018 · 2a14998 · 2a14998
1 parent 0a81ea7
commit 2a14998
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 13 deletions.
diff --git a/app/controllers/barbeque/sns_subscriptions_controller.rb b/app/controllers/barbeque/sns_subscriptions_controller.rb
@@ -1,5 +1,3 @@
-require 'aws-sdk-sns'
-
 class Barbeque::SnsSubscriptionsController < Barbeque::ApplicationController
   def index
     @sns_subscriptions = Barbeque::SNSSubscription.all
@@ -47,10 +45,10 @@ def destroy
 
   def fetch_sns_topic_arns
     if Barbeque.config.sns_regions.empty?
-      Aws::SNS::Client.new.list_topics.topics.map(&:topic_arn)
+      Barbeque::SNSSubscriptionService.sns_client.list_topics.topics.map(&:topic_arn)
     else
       Barbeque.config.sns_regions.flat_map do |region|
-        Aws::SNS::Client.new(region: region).list_topics.topics.map(&:topic_arn)
+        Barbeque::SNSSubscriptionService.sns_client(region).list_topics.topics.map(&:topic_arn)
       end
     end
   end

diff --git a/app/services/barbeque/sns_subscription_service.rb b/app/services/barbeque/sns_subscription_service.rb
@@ -6,6 +6,11 @@ def self.sqs_client
     @sqs_client ||= Aws::SQS::Client.new
   end
 
+  def self.sns_client(region = nil)
+    @sns_clients ||= {}
+    @sns_clients[region] ||= Aws::SNS::Client.new(region: region)
+  end
+
   # @param [Barbeque::SNSSubscription] sns_subscription
   # @return [Boolean] `true` if succeeded to subscribe
   def subscribe(sns_subscription)
@@ -38,7 +43,11 @@ def unsubscribe(sns_subscription)
   private
 
   def sqs_client
-    Barbeque::SNSSubscriptionService.sqs_client
+    self.class.sqs_client
+  end
+
+  def sns_client(region)
+    self.class.sns_client(region)
   end
 
   # @param [Barbeque::SNSSubscription] sns_subscription
@@ -90,8 +99,7 @@ def subscribe_topic!(sns_subscription)
     )
     queue_arn = sqs_attrs.attributes['QueueArn']
 
-    sns_client = Aws::SNS::Client.new(region: sns_subscription.region)
-    sns_client.subscribe(
+    sns_client(sns_subscription.region).subscribe(
       topic_arn: sns_subscription.topic_arn,
       protocol: 'sqs',
       endpoint: queue_arn
@@ -106,15 +114,13 @@ def unsubscribe_topic!(sns_subscription)
     )
     queue_arn = sqs_attrs.attributes['QueueArn']
 
-    sns_client = Aws::SNS::Client.new(region: sns_subscription.region)
-
-    subscriptions = sns_client.list_subscriptions_by_topic(
+    subscriptions = sns_client(sns_subscription.region).list_subscriptions_by_topic(
       topic_arn: sns_subscription.topic_arn,
     )
     subscription_arn = subscriptions.subscriptions.find {|subscription| subscription.endpoint == queue_arn }.try!(:subscription_arn)
 
     if subscription_arn
-      sns_client.unsubscribe(
+      sns_client(sns_subscription.region).unsubscribe(
         subscription_arn: subscription_arn,
       )
     end

diff --git a/spec/controllers/barbeque/job_definitions_controller_spec.rb b/spec/controllers/barbeque/job_definitions_controller_spec.rb
@@ -191,8 +191,8 @@
       let(:subscription_arn) { 'arn:aws:sns:ap-northeast-1:012345678912:barbeque-spec:01234567-89ab-cdef-0123-456789abcdef' }
 
       before do
+        allow(Barbeque::SNSSubscriptionService).to receive(:sns_client).with('ap-northeast-1').and_return(sns_client)
         allow(Barbeque::SNSSubscriptionService).to receive(:sqs_client).and_return(sqs_client)
-        allow(Aws::SNS::Client).to receive(:new).with(region: 'ap-northeast-1').and_return(sns_client)
 
         allow(sqs_client).to receive(:get_queue_attributes).
           with(queue_url: sns_subscription.job_queue.queue_url, attribute_names: ['QueueArn']).

diff --git a/spec/controllers/barbeque/sns_subscriptions_controller_spec.rb b/spec/controllers/barbeque/sns_subscriptions_controller_spec.rb
@@ -5,7 +5,7 @@
 
   before do
     allow(Barbeque::SNSSubscriptionService).to receive(:sqs_client).and_return(sqs_client)
-    allow(Aws::SNS::Client).to receive(:new).with(region: 'ap-northeast-1').and_return(sns_client)
+    allow(Barbeque::SNSSubscriptionService).to receive(:sns_client).with('ap-northeast-1').and_return(sns_client)
   end
 
   describe '#create' do