lib/td/logger/agent/rails.rb in td-logger-0.2.0 vs lib/td/logger/agent/rails.rb in td-logger-0.2.1

@@ -26,10 +26,26 @@
 
 # disable logging
 test:
 EOF
 
+    def self.rails_env
+      if defined?(RAILS_ENV) && RAILS_ENV.to_s != ''
+        RAILS_ENV.to_s
+      else
+        ::Rails.env
+      end
+    end
+    
+    def self.rails_root
+      if defined?(RAILS_ROOT) && RAILS_ROOT.to_s != ''
+        RAILS_ROOT.to_s
+      else
+        ::Rails.root
+      end
+    end
+
     class Config
       def initialize(conf)
         if agent = conf['agent']
           host, port = agent.split(':',2)
           port = (port || 24224).to_i
@@ -67,21 +83,21 @@
     end
 
     def self.read_config(rails)
       logger = ::Rails.logger || ::Logger.new(STDOUT)
       begin
-        yaml = YAML.load_file("#{RAILS_ROOT}/#{CONFIG_PATH}")
+        yaml = YAML.load_file("#{rails_root}/#{CONFIG_PATH}")
       rescue
         logger.warn "Can't load #{CONFIG_PATH} file."
         logger.warn "  #{$!}"
         logger.warn "Put the following file:"
-        logger.warn sample
+        logger.warn CONFIG_SAMPLE
         return
       end
 
-      conf = yaml[RAILS_ENV]
+      conf = yaml[rails_env]
       unless conf
-        logger.warn "#{CONFIG_PATH} doesn't include setting for current environment (#{RAILS_ENV})."
+        logger.warn "#{CONFIG_PATH} doesn't include setting for current environment (#{rails_env})."
         logger.warn "Disabling Treasure Data logger."
         return
       end
 
       begin