metasploit-framework/spec/support/shared/contexts/database_cleaner.rb

# -*- coding:binary -*-
require 'metasploit/framework/database'

shared_context 'DatabaseCleaner' do
	def with_established_connection
		begin
			ActiveRecord::Base.connection_pool.with_connection do
				yield
			end
		rescue ActiveRecord::ConnectionNotEstablished
			# if there isn't a connection established, then established one and try
			# again
			ActiveRecord::Base.configurations = Metasploit::Framework::Database.configurations
			spec = ActiveRecord::Base.configurations[Metasploit::Framework.env]
			ActiveRecord::Base.establish_connection(spec)

			retry
		end
	end

	# clean before all in case last test run was interrupted before
	# after(:each) could clean up
	before(:all) do
		with_established_connection do
			DatabaseCleaner.clean_with(:truncation)
		end
	end

	# Clean up after each test
	after(:each) do
		with_established_connection do
			# Testing using both :truncation and :deletion; :truncation took long
			# for testing.
			DatabaseCleaner.clean_with(:deletion)
		end
	end
end