lib/test_queue/iterator.rb in test-queue-0.1.3 vs lib/test_queue/iterator.rb in test-queue-0.2.0.beta.1

@@ -1,41 +1,64 @@
 module TestQueue
   class Iterator
-    attr_reader :stats
+    attr_reader :stats, :sock
 
-    def initialize(sock)
-      @sock = sock
+    def initialize(sock, suites, filter=nil)
       @done = false
       @stats = {}
       @procline = $0
+      @sock = sock
+      @suites = suites
+      @filter = filter
+      if @sock =~ /^(.+):(\d+)$/
+        @tcp_address = $1
+        @tcp_port = $2.to_i
+      end
     end
 
     def each
       fail 'already used this iterator' if @done
 
       while true
-        client = UNIXSocket.new(@sock)
+        client = connect_to_master('POP')
         r, w, e = IO.select([client], nil, [client], nil)
         break if !e.empty?
 
-        if data = client.read(16384)
+        if data = client.read(65536)
           client.close
           item = Marshal.load(data)
-          $0 = "#{@procline} - #{item.respond_to?(:description) ? item.description : item}"
+          break if item.nil?
+          suite = @suites[item]
 
+          $0 = "#{@procline} - #{suite.respond_to?(:description) ? suite.description : suite}"
           start = Time.now
-          yield item
-          @stats[item] = Time.now - start
+          if @filter
+            @filter.call(suite){ yield suite }
+          else
+            yield suite
+          end
+          @stats[suite.to_s] = Time.now - start
         else
           break
         end
       end
     rescue Errno::ENOENT, Errno::ECONNRESET, Errno::ECONNREFUSED
     ensure
       @done = true
       File.open("/tmp/test_queue_worker_#{$$}_stats", "wb") do |f|
         f.write Marshal.dump(@stats)
       end
+    end
+
+    def connect_to_master(cmd)
+      sock =
+        if @tcp_address
+          TCPSocket.new(@tcp_address, @tcp_port)
+        else
+          UNIXSocket.new(@sock)
+        end
+      sock.puts(cmd)
+      sock
     end
 
     include Enumerable
 
     def empty?