spec/test_server.rb in instrumental_agent-0.8.1 vs spec/test_server.rb in instrumental_agent-0.8.2

@@ -1,14 +1,21 @@
 class TestServer
   attr_accessor :host, :port, :connect_count, :commands
 
-  def initialize
+  def initialize(options={})
+    default_options = {
+      :listen => true,
+      :authenticate => true,
+      :response => true,
+    }
+    @options = default_options.merge(options)
+
     @connect_count = 0
     @connections = []
     @commands = []
     @host = 'localhost'
-    listen
+    listen if @options[:listen]
   end
 
   def listen
     @port ||= 10001
     @server = TCPServer.new(port)
@@ -23,10 +30,19 @@
             # puts "connection received"
             loop do
               command = socket.gets.strip
               # puts "got: #{command}"
               commands << command
+              if %w[hello authenticate].include?(command.split(' ')[0])
+                if @options[:response]
+                  if @options[:authenticate]
+                    socket.puts "ok"
+                  else
+                    socket.puts "gtfo"
+                  end
+                end
+              end
             end
           end
         end
       rescue Exception => err
         unless @stopping
@@ -48,10 +64,10 @@
   end
 
   def stop
     @stopping = true
     disconnect_all
-    @server.close # FIXME: necessary?
+    @server.close if @server
   end
 
   def disconnect_all
     @connections.each { |c| c.close rescue false }
   end