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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198
|
require "wait_for_it/version"
require 'pathname'
require 'shellwords'
require 'tempfile'
require 'timeout'
class WaitForIt
class WaitForItTimeoutError < StandardError
def initialize(options = {})
command = options[:command]
input = options[:input]
timeout = options[:timeout]
log = options[:log]
super "Running command: '#{ command }', waiting for '#{ input }' did not occur within #{ timeout } seconds:\n#{ log.read }"
end
end
DEFAULT_TIMEOUT = 10 # seconds
DEFAULT_OUT = ">>"
DEFAULT_ENV = {}
@wait_for = nil
@timeout = nil
@redirection = nil
@env = nil
# Configure global WaitForIt settings
def self.config
yield self
self
end
# The default output is expected in the logs before the process is considered "booted"
def self.wait_for=(wait_for)
@wait_for = wait_for
end
def self.wait_for
@wait_for
end
# The default timeout that is waited for a process to boot
def self.timeout=(timeout)
@timeout = timeout
end
def self.timeout
@timeout || DEFAULT_TIMEOUT
end
# The default shell redirect to the logs
def self.redirection=(redirection)
@redirection = redirection
end
def self.redirection
@redirection || DEFAULT_OUT
end
# Default environment variables under which commands should be executed.
def self.env=(env)
@env = env
end
def self.env
@env || DEFAULT_ENV
end
# Creates a new WaitForIt instance
#
# @param [String] command Command to spawn
# @param [Hash] options
# @options options [Fixnum] :timeout The duration to wait a commmand to boot, default is 10 seconds
# @options options [String] :wait_for The output the process emits when it has successfully booted.
# When present the calling process will block until the message is received in the log output
# or until the timeout is hit.
# @options options [String] :redirection The shell redirection used to pipe to log file
# @options options [Hash] :env Keys and values for environment variables in the process
def initialize(command, options = {})
@command = command
@timeout = options[:timeout] || WaitForIt.timeout
@wait_for = options[:wait_for] || WaitForIt.wait_for
redirection = options[:redirection] || WaitForIt.redirection
env = options[:env] || WaitForIt.env
@log = set_log
@pid = nil
raise "Must provide a wait_for: option" unless @wait_for
spawn(command, redirection, env)
wait!(@wait_for)
if block_given?
begin
yield self
ensure
cleanup
end
end
rescue WaitForItTimeoutError => e
cleanup
raise e
end
attr_reader :timeout, :log
# Checks the logs of the process to see if they contain a match.
# Can use a string or a regular expression.
def contains?(input)
log.read.match convert_to_regex(input)
end
# Returns a count of the number of times logs match the input.
# Can use a string or a regular expression.
def count(input)
log.read.scan(convert_to_regex(input)).count
end
# Blocks parent process until given message appears at the
def wait(input, t = timeout)
regex = convert_to_regex(input)
Timeout::timeout(t) do
until log.read.match regex
sleep 0.01
end
end
sleep 0.01
self
rescue Timeout::Error
puts "Timeout waiting for #{input.inspect} to find a match using #{ regex } in \n'#{ log.read }'"
false
end
# Same as `wait` but raises an error if timeout is reached
def wait!(input, t = timeout)
unless wait(input, t)
options = {}
options[:command] = @command
options[:input] = input
options[:timeout] = t
options[:log] = @log
raise WaitForItTimeoutError.new(options)
end
end
# Kills the process and removes temporary files
def cleanup
shutdown
close_log
unlink_log
end
private
def close_log
@tmp_file.close if @tmp_file
end
def unlink_log
@log.unlink if @log
rescue Errno::ENOENT
# File already unlinked
end
def set_log
@tmp_file = Tempfile.new(["wait_for_it", ".log"])
log_file = Pathname.new(@tmp_file)
log_file.mkpath unless log_file.exist?
log_file
end
def spawn(command, redirection, env_hash = {})
env = {}
env_hash.each {|k, v| env[k.to_s] = v.nil? ? v : v.to_s }
# Must exec so when we kill the PID it kills the child process
@pid = Process.spawn(env, "exec #{ command } #{ redirection } #{ log }")
end
def convert_to_regex(input)
return input if input.is_a?(Regexp)
Regexp.new(Regexp.escape(input))
end
# Kills the process and waits for it to exit
def shutdown
if @pid
Process.kill('TERM', @pid)
Process.wait(@pid)
@pid = nil
end
rescue Errno::ESRCH
# Process doesn't exist, nothing to kill
end
end
|