lib/unleash/configuration.rb in unleash-4.2.1 vs lib/unleash/configuration.rb in unleash-4.3.0

@@ -21,11 +21,11 @@
       :logger,
       :log_level,
       :bootstrap_config
 
     def initialize(opts = {})
-      ensure_valid_opts(opts)
+      validate_custom_http_headers!(opts[:custom_http_headers]) if opts.has_key?(:custom_http_headers)
       set_defaults
 
       initialize_default_logger if opts[:logger].nil?
 
       merge(opts)
@@ -38,22 +38,23 @@
 
     def validate!
       return if self.disable_client
 
       raise ArgumentError, "URL and app_name are required parameters." if self.app_name.nil? || self.url.nil?
-      raise ArgumentError, "custom_http_headers must be a hash." unless self.custom_http_headers.is_a?(Hash)
+
+      validate_custom_http_headers!(self.custom_http_headers)
     end
 
     def refresh_backup_file!
       self.backup_file = File.join(Dir.tmpdir, "unleash-#{app_name}-repo.json")
     end
 
     def http_headers
       {
         'UNLEASH-INSTANCEID' => self.instance_id,
         'UNLEASH-APPNAME' => self.app_name
-      }.merge(custom_http_headers.dup)
+      }.merge!(generate_custom_http_headers)
     end
 
     def fetch_toggles_uri
       uri = URI("#{self.url_stripped_of_slash}/client/features")
       uri.query = "project=#{self.project_name}" unless self.project_name.nil?
@@ -76,16 +77,10 @@
       self.bootstrap_config&.valid?
     end
 
     private
 
-    def ensure_valid_opts(opts)
-      unless opts[:custom_http_headers].is_a?(Hash) || opts[:custom_http_headers].nil?
-        raise ArgumentError, "custom_http_headers must be a hash."
-      end
-    end
-
     def set_defaults
       self.app_name         = nil
       self.environment      = 'default'
       self.url              = nil
       self.instance_id      = SecureRandom.uuid
@@ -114,9 +109,21 @@
     end
 
     def merge(opts)
       opts.each_pair{ |opt, val| set_option(opt, val) }
       self
+    end
+
+    def validate_custom_http_headers!(custom_http_headers)
+      return if custom_http_headers.is_a?(Hash) || custom_http_headers.respond_to?(:call)
+
+      raise ArgumentError, "custom_http_headers must be a Hash or a Proc."
+    end
+
+    def generate_custom_http_headers
+      return self.custom_http_headers.call if self.custom_http_headers.respond_to?(:call)
+
+      self.custom_http_headers
     end
 
     def set_option(opt, val)
       __send__("#{opt}=", val)
     rescue NoMethodError