diff options
| -rw-r--r-- | spec/util_queue_spec.lua | 37 | ||||
| -rw-r--r-- | util/queue.lua | 3 |
2 files changed, 40 insertions, 0 deletions
diff --git a/spec/util_queue_spec.lua b/spec/util_queue_spec.lua index 7cd3d695..d73f523d 100644 --- a/spec/util_queue_spec.lua +++ b/spec/util_queue_spec.lua @@ -100,4 +100,41 @@ describe("util.queue", function() end); end); + describe("consume()", function () + it("should work", function () + local q = queue.new(10); + for i = 1, 5 do + q:push(i); + end + local c = 0; + for i in q:consume() do + assert(i == c + 1); + assert(q:count() == (5-i)); + c = i; + end + end); + + it("should work even if items are pushed in the loop", function () + local q = queue.new(10); + for i = 1, 5 do + q:push(i); + end + local c = 0; + for i in q:consume() do + assert(i == c + 1); + if c < 3 then + assert(q:count() == (5-i)); + else + assert(q:count() == (6-i)); + end + + c = i; + + if c == 3 then + q:push(6); + end + end + assert.equal(c, 6); + end); + end); end); diff --git a/util/queue.lua b/util/queue.lua index 728e905f..e63b3f1c 100644 --- a/util/queue.lua +++ b/util/queue.lua @@ -64,6 +64,9 @@ local function new(size, allow_wrapping) return pos+1, t._items[read_pos]; end, self, 0; end; + consume = function (self) + return self.pop, self; + end; }; end |