/
base.rb
75 lines (58 loc) · 1.98 KB
/
base.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
require 'database_cleaner/generic/base'
require 'active_record'
require 'erb'
module DatabaseCleaner
module ActiveRecord
def self.available_strategies
%w[truncation transaction deletion]
end
def self.config_file_location=(path)
@config_file_location = path
end
def self.config_file_location
@config_file_location ||= "#{DatabaseCleaner.app_root}/config/database.yml"
end
module Base
include ::DatabaseCleaner::Generic::Base
attr_accessor :connection_hash
def db=(desired_db)
@db = desired_db
load_config
end
def db
@db || super
end
def load_config
if self.db != :default && File.file?(ActiveRecord.config_file_location)
connection_details = YAML::load(ERB.new(IO.read(ActiveRecord.config_file_location)).result)
@connection_hash = connection_details[self.db.to_s]
end
end
def create_connection_klass
Class.new(::ActiveRecord::Base)
end
def connection_klass
@connection_klass ||= if @db == :default || (@db.nil? && connection_hash.nil?)
::ActiveRecord::Base
elsif connection_hash
lookup_from_connection_pool || establish_connection
else
@db # allows for an actual class to be passed in
end
end
private
def lookup_from_connection_pool
if ::ActiveRecord::Base.respond_to?(:descendants)
database_name = connection_hash["database"] || connection_hash[:database]
models = ::ActiveRecord::Base.descendants
models.detect {|m| m.connection_pool.spec.config[:database] == database_name}
end
end
def establish_connection
klass = create_connection_klass
klass.send :establish_connection, connection_hash
klass
end
end
end
end