add safeguard for running in production, namespace errors, add error messages

lib/database_cleaner/safeguard.rb

@@ -1,21 +1,70 @@
 module DatabaseCleaner
-  DatabaseUrlSpecified = Class.new(Exception)
-
   class Safeguard
-    def run
+    class Error < Exception
-      return if skip?
+      class RemoteDatabaseUrl < Error
-      raise DatabaseUrlSpecified if env_db_url?
+        def initialize
+          super("ENV['DATABASE_URL'] is set to a remote URL. Please refer to https://github.com/DatabaseCleaner/database_cleaner#safeguards")
+        end
+      end
+
+      class ProductionEnv < Error
+        def initialize(env)
+          super("ENV['#{env}'] is set to production. Please refer to https://github.com/DatabaseCleaner/database_cleaner#safeguards")
+        end
+      end
     end
 
-    private
+    class RemoteDatabaseUrl
-
+      def run
-      def env_db_url?
+        raise Error::RemoteDatabaseUrl if !skip? && given?
-        url = ENV['DATABASE_URL']
-        url && !url.include?('localhost')
       end
 
-      def skip?
+      private
-        !!ENV['DATABASE_CLEANER_SKIP_SAFEGUARD']
+
+        def given?
+          remote?(ENV['DATABASE_URL'])
+        end
+
+        def remote?(url)
+          url && !url.include?('localhost')
+        end
+
+        def skip?
+          ENV['DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL'] ||
+            DatabaseCleaner.allow_remote_database_url
+        end
+    end
+
+    class Production
+      KEYS = %w(ENV RACK_ENV RAILS_ENV)
+
+      def run
+        raise Error::ProductionEnv.new(key) if !skip? && given?
       end
+
+      private
+
+        def given?
+          !!key
+        end
+
+        def key
+          @key ||= KEYS.detect { |key| ENV[key] == 'production' }
+        end
+
+        def skip?
+          ENV['DATABASE_CLEANER_ALLOW_PRODUCTION'] ||
+            DatabaseCleaner.allow_production
+        end
+    end
+
+    CHECKS = [
+      RemoteDatabaseUrl,
+      Production
+    ]
+
+    def run
+      CHECKS.each { |const| const.new.run }
+    end
   end
 end

spec/database_cleaner/safeguard_spec.rb

@@ -1,9 +1,12 @@
 require 'spec_helper'
+require 'support/env'
 require 'active_record'
 require 'database_cleaner/active_record/transaction'
 
 module DatabaseCleaner
   describe Safeguard do
+    include Support::Env
+
     let(:strategy) { DatabaseCleaner::ActiveRecord::Transaction }
     let(:cleaner)  { Base.new(:autodetect) }
 
@@ -11,31 +14,66 @@ module DatabaseCleaner
 
     describe 'DATABASE_URL is set' do
       describe 'to any value' do
-        before { ENV['DATABASE_URL'] = 'postgres://remote.host' }
+        env DATABASE_URL: 'postgres://remote.host'
-        after  { ENV.delete('DATABASE_URL') }
 
-        it 'raises DatabaseUrlSpecified' do
+        it 'raises' do
-          expect { cleaner.start }.to raise_error(DatabaseUrlSpecified)
+          expect { cleaner.start }.to raise_error(Safeguard::Error::RemoteDatabaseUrl)
         end
       end
 
       describe 'to a local url' do
-        before { ENV['DATABASE_URL'] = 'postgres://localhost' }
+        env DATABASE_URL: 'postgres://localhost'
-        after  { ENV.delete('DATABASE_URL') }
 
         it 'does not raise' do
           expect { cleaner.start }.to_not raise_error
         end
       end
 
-      describe 'DATABASE_CLEANER_SKIP_SAFEGUARD is set' do
+      describe 'DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL is set' do
-        before { ENV['DATABASE_CLEANER_SKIP_SAFEGUARD'] = 'true' }
+        env DATABASE_CLEANER_ALLOW_REMOTE_DATABASE_URL: true
-        after  { ENV.delete('DATABASE_CLEANER_SKIP_SAFEGUARD') }
+
+        it 'does not raise' do
+          expect { cleaner.start }.to_not raise_error
+        end
+      end
+
+      describe 'DatabaseCleaner.allow_remote_database_url is true' do
+        before { DatabaseCleaner.allow_remote_database_url = true }
+        after  { DatabaseCleaner.allow_remote_database_url = nil }
 
         it 'does not raise' do
           expect { cleaner.start }.to_not raise_error
         end
       end
     end
+
+    describe 'ENV is set to production' do
+      %w(ENV RACK_ENV RAILS_ENV).each do |key|
+        describe "on #{key}" do
+          env key => 'production'
+
+          it 'raises' do
+            expect { cleaner.start }.to raise_error(Safeguard::Error::ProductionEnv)
+          end
+        end
+
+        describe 'DATABASE_CLEANER_ALLOW_PRODUCTION is set' do
+          env DATABASE_CLEANER_ALLOW_PRODUCTION: true
+
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
+        end
+
+        describe 'DatabaseCleaner.allow_production is true' do
+          before { DatabaseCleaner.allow_production = true }
+          after  { DatabaseCleaner.allow_production = nil }
+
+          it 'does not raise' do
+            expect { cleaner.start }.to_not raise_error
+          end
+        end
+      end
+    end
   end
 end

spec/support/env.rb

@@ -0,0 +1,22 @@
+module Support
+  module Env
+    def self.included(base)
+      base.send(:extend, ClassMethods)
+    end
+
+    module ClassMethods
+      def env(vars)
+        before { define_env(vars) }
+        after  { undefine_env(vars) }
+      end
+    end
+
+    def define_env(vars)
+      vars.each { |key, value| ENV[key.to_s.upcase] = value.to_s }
+    end
+
+    def undefine_env(vars)
+      vars.each { |key, _| ENV.delete(key.to_s) }
+    end
+  end
+end