From 8998fe906d0d8673164deb3be267f82a35a70659 Mon Sep 17 00:00:00 2001
From: Supreet <supreet@kikoff.com>
Date: Fri, 25 Nov 2022 15:39:29 -0500
Subject: [PATCH] Add timeout options to the config

---
 lib/amplitude_api.rb           | 23 ++++++++++++++++++++---
 lib/amplitude_api/config.rb    |  7 +++++--
 spec/lib/amplitude_api_spec.rb |  3 ++-
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/lib/amplitude_api.rb b/lib/amplitude_api.rb
index 8fc285f..4d1fc26 100644
--- a/lib/amplitude_api.rb
+++ b/lib/amplitude_api.rb
@@ -89,7 +89,11 @@ def track_body(*events)
     #
     # Send one or more Events to the Amplitude API
     def track(*events)
-      Faraday.post(TRACK_URI_STRING, track_body(events), "Content-Type" => "application/json")
+      Faraday.post(TRACK_URI_STRING, track_body(events), "Content-Type" => "application/json") do |request|
+        request.options.open_timeout = config.open_timeout
+        request.options.read_timeout = config.read_timeout
+        request.options.write_timeout = config.write_timeout
+      end
     end
 
     # ==== Identification related methods
@@ -132,7 +136,11 @@ def identify_body(*identifications)
     #
     # Send one or more Identifications to the Amplitude Identify API
     def identify(*identifications)
-      Faraday.post(IDENTIFY_URI_STRING, identify_body(identifications))
+      Faraday.post(IDENTIFY_URI_STRING, identify_body(identifications)) do |request|
+        request.options.open_timeout = config.open_timeout
+        request.options.read_timeout = config.read_timeout
+        request.options.write_timeout = config.write_timeout
+      end
     end
 
     # ==== Event Segmentation related methods
@@ -157,7 +165,7 @@ def identify(*identifications)
     #
     # @return [ Faraday::Response ]
     def segmentation(event, start_time, end_time, **options)
-      Faraday.get SEGMENTATION_URI_STRING, userpwd: "#{api_key}:#{secret_key}", params: {
+      params = {
         e: event.to_json,
         m: options[:m],
         start: start_time.strftime("%Y%m%d"),
@@ -167,6 +175,12 @@ def segmentation(event, start_time, end_time, **options)
         g: options[:g],
         limit: options[:limit]
       }.delete_if { |_, value| value.nil? }
+
+      Faraday.get(SEGMENTATION_URI_STRING, userpwd: "#{api_key}:#{secret_key}", params: params) do |request|
+        request.options.open_timeout = config.open_timeout
+        request.options.read_timeout = config.read_timeout
+        request.options.write_timeout = config.write_timeout
+      end
     end
 
     # Delete a user from amplitude
@@ -190,6 +204,9 @@ def delete(user_ids: nil, amplitude_ids: nil, requester: nil, ignore_invalid_id:
 
       faraday = Faraday.new do |conn|
         conn.request :basic_auth, config.api_key, config.secret_key
+        conn.options.open_timeout = config.open_timeout
+        conn.options.read_timeout = config.read_timeout
+        conn.options.write_timeout = config.write_timeout
       end
 
       faraday.post(
diff --git a/lib/amplitude_api/config.rb b/lib/amplitude_api/config.rb
index 3a9613d..2c30fd0 100644
--- a/lib/amplitude_api/config.rb
+++ b/lib/amplitude_api/config.rb
@@ -9,7 +9,7 @@ class Config
 
     attr_accessor :api_key, :secret_key, :whitelist, :time_formatter,
                   :event_properties_formatter, :user_properties_formatter,
-                  :options
+                  :options, :timeout, :open_timeout, :read_timeout, :write_timeout
 
     def initialize
       self.class.defaults.each { |k, v| send("#{k}=", v) }
@@ -44,7 +44,10 @@ def defaults
           whitelist: base_properties + revenue_properties + optional_properties,
           time_formatter: ->(time) { time ? time.to_i * 1_000 : nil },
           event_properties_formatter: ->(props) { props || {} },
-          user_properties_formatter: ->(props) { props || {} }
+          user_properties_formatter: ->(props) { props || {} },
+          open_timeout: 60, # Net::HTTP default. Number of seconds to wait for the connection to open
+          read_timeout: 60, # Net::HTTP default. Number of seconds to wait for one block to be read
+          write_timeout: 60 # Net::HTTP default. Number of seconds to wait for one block to be written
         }
       end
     end
diff --git a/spec/lib/amplitude_api_spec.rb b/spec/lib/amplitude_api_spec.rb
index 9770657..62e824b 100644
--- a/spec/lib/amplitude_api_spec.rb
+++ b/spec/lib/amplitude_api_spec.rb
@@ -432,7 +432,8 @@
   end
 
   describe ".delete" do
-    let(:connection) { instance_double("Faraday::Connection", post: nil, request: nil) }
+    let(:connection) { instance_double("Faraday::Connection", post: nil, request: nil, options: options) }
+    let(:options) { Faraday::RequestOptions.new }
 
     before do
       allow(Faraday).to receive(:new).and_yield(connection).and_return(connection)