File: wait_for_it.rb

package info (click to toggle)
ruby-wait-for-it 0.2.1-2
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye, forky, sid, trixie
  • size: 124 kB
  • sloc: ruby: 146; sh: 4; makefile: 4
file content (198 lines) | stat: -rw-r--r-- 4,908 bytes parent folder | download
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