From 8d2ea04c4f24253c861cd1fff1e0cbe4e2a3825a Mon Sep 17 00:00:00 2001
From: Holger Just <hello@holgerjust.de>
Date: Fri, 3 Feb 2017 00:28:45 +0100
Subject: [PATCH] Set a timestamp on the Buffer

---
 lib/rackstash/buffer.rb       | 21 +++++++++++++++
 spec/rackstash/buffer_spec.rb | 50 ++++++++++++++++++++++++++++++++---
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/lib/rackstash/buffer.rb b/lib/rackstash/buffer.rb
index 7cc65b0..0da76e1 100644
--- a/lib/rackstash/buffer.rb
+++ b/lib/rackstash/buffer.rb
@@ -40,6 +40,7 @@ module Rackstash
       @messages = []
       @fields = Rackstash::Fields::Hash.new(forbidden_keys: FORBIDDEN_FIELDS)
       @tags = Rackstash::Fields::Tags.new
+      @timestamp = nil
     end
 
     # Add a new message to the buffer. This will mark the current buffer as
@@ -50,6 +51,9 @@ module Rackstash
     # @return [Message] the passed `message`
     def add_message(message)
       @messages << message
+      timestamp(message.time)
+
+      message
     end
 
     def allow_empty?
@@ -103,5 +107,22 @@ module Rackstash
     def tag(*tags, scope: nil)
       @tags.merge!(tags, scope: scope)
     end
+
+    # Returns the time of the current buffer as an ISO 8601 formatted string.
+    # If the timestamp was not yet set on the buffer, it is is set to the
+    # the passed `time` or the current time.
+    #
+    # @example
+    #   buffer.timestamp
+    #   # => "2016-10-17T13:37:00.234Z"
+    # @param time [Time] an optional time object. If no timestamp was set yet,
+    #   this time is used
+    # @return [String] an ISO 8601 formatted UTC timestamp.
+    def timestamp(time = nil)
+      @timestamp ||= begin
+        time ||= Time.now
+        time.utc.iso8601(ISO8601_PRECISION).freeze
+      end
+    end
   end
 end
diff --git a/spec/rackstash/buffer_spec.rb b/spec/rackstash/buffer_spec.rb
index a4518bc..8887de7 100644
--- a/spec/rackstash/buffer_spec.rb
+++ b/spec/rackstash/buffer_spec.rb
@@ -19,16 +19,24 @@ describe Rackstash::Buffer do
 
   describe '#add_message' do
     it 'adds a message to the buffer' do
-      msg = double(message: 'Hello World')
-      buffer.add_message msg
+      msg = double(message: 'Hello World', time: Time.now)
+      expect(buffer.add_message msg).to equal msg
 
       expect(buffer.messages).to eql [msg]
     end
+
+    it 'sets the timestamp' do
+      time = Time.parse('2016-10-17 13:37:00 +03:00')
+      msg = double(message: 'Hello World', time: time)
+
+      buffer.add_message msg
+      expect(buffer.timestamp).to eql '2016-10-17T10:37:00.000Z'
+    end
   end
 
   describe '#messages' do
     it 'returns an array of messages' do
-      msg = double('Rackstash::Message')
+      msg = double(message: 'Hello World', time: Time.now)
       buffer.add_message(msg)
 
       expect(buffer.messages).to eql [msg]
@@ -45,7 +53,7 @@ describe Rackstash::Buffer do
 
   describe '#pending?' do
     it 'sets pending when adding a message' do
-      buffer.add_message double(message: 'some message')
+      buffer.add_message double(message: 'some message', time: Time.now)
       expect(buffer.pending?).to be true
     end
 
@@ -140,4 +148,38 @@ describe Rackstash::Buffer do
       end
     end
   end
+
+  describe '#timestamp' do
+    it 'initializes @timestamp to Time.now.utc' do
+      now = Time.parse('2016-10-17 13:37:00 +03:00')
+
+      expect(Time).to receive(:now).once.and_return(now)
+      expect(now).to receive(:utc).once.and_return(now.utc)
+
+      expect(buffer.timestamp).to eql '2016-10-17T10:37:00.000Z'
+      expect(buffer.timestamp).to eql '2016-10-17T10:37:00.000Z'
+    end
+
+    it 'initializes @timestamp with the passed time' do
+      now = Time.parse('2016-10-17 13:37:00 +03:00')
+
+      expect(Time).not_to receive(:now)
+      expect(buffer.timestamp(now)).to eql '2016-10-17T10:37:00.000Z'
+      expect(buffer.timestamp).to eql '2016-10-17T10:37:00.000Z'
+    end
+
+    it 'does not overwrites an already set timestamp' do
+      first = Time.parse('2016-10-17 10:10:10 +03:00')
+      second = Time.parse('2016-10-17 20:20:20 +03:00')
+
+      buffer.timestamp(first)
+      expect(buffer.timestamp).to eql '2016-10-17T07:10:10.000Z'
+
+      buffer.timestamp
+      expect(buffer.timestamp).to eql '2016-10-17T07:10:10.000Z'
+
+      buffer.timestamp(second)
+      expect(buffer.timestamp).to eql '2016-10-17T07:10:10.000Z'
+    end
+  end
 end