lib/rack/heroku_env.rb in rack-heroku_env-0.0.1 vs lib/rack/heroku_env.rb in rack-heroku_env-0.1.0

@@ -1,26 +1,31 @@
 module Rack
   class HerokuEnv
-    VERSION = "0.0.1"
+    VERSION = '0.1.0'
 
-    VARIABLE_NAMES = {
-      'HTTP_X_HEROKU_DYNOS_IN_USE' => 'X-Heroku-Dynos-In-Use',
-      'HTTP_X_HEROKU_QUEUE_DEPTH' => 'X-Heroku-Queue-Depth',
-      'HTTP_X_HEROKU_QUEUE_WAIT_TIME' => 'X-Heroku-Queue-Wait-Time',
-    }
-
     def initialize(app)
       @app = app
     end
 
     def call(env)
       response = @app.call(env)
 
-      VARIABLE_NAMES.each do |variable_name, header_name|
-        response[1][header_name] = env[variable_name] if env[variable_name]
+      heroku_vars(env).each do |variable_name|
+        response[1][header_name_for(variable_name)] = env[variable_name]
       end
 
       response
+    end
+
+    private
+
+
+    def heroku_vars(env)
+      env.keys.grep(/\AHTTP_X_HEROKU/)
+    end
+
+    def header_name_for(variable_name)
+      variable_name.gsub(/\AHTTP_/, '').split('_').map{|w| w.capitalize }.join('-')
     end
 
   end
 end