From 6289baaa503b6c80893dfc2a6efae80a35a6764c Mon Sep 17 00:00:00 2001
From: contradict <contradict@gmail.com>
Date: Tue, 7 Sep 2021 16:02:11 -0700
Subject: [PATCH] Add more iteration tests

---
 test/runtests.jl | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index b14b36e..d4bedbf 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -269,8 +269,7 @@ const encodings = [
     @test YAML.load(IOBuffer(data)) == "test"
 end
 
-@testset "multi_doc_bom" begin
-    iterable = YAML.load_all("""
+const multidoc_contents = """
 \ufeff---\r
 test: 1
 \ufeff---
@@ -278,7 +277,25 @@ test: 2
 
 \ufeff---
 test: 3
-""")
+"""
+
+@testset "multi_doc_bom" begin
+    iterable = YAML.load_all(multidoc_contents)
+    (val, state) = iterate(iterable)
+    @test equivalent(val, Dict("test" => 1))
+    (val, state) = iterate(iterable, state)
+    @test equivalent(val, Dict("test" => 2))
+    (val, state) = iterate(iterable, state)
+    @test equivalent(val, Dict("test" => 3))
+    @test iterate(iterable, state) === nothing
+end
+
+@testset "multi_doc_file" begin
+    fname = tempname() # cleanup=true, file will be deleted on process exit
+    open(fname, "w") do f
+        write(f, multidoc_contents)
+    end
+    iterable = YAML.load_all_file(fname)
     (val, state) = iterate(iterable)
     @test equivalent(val, Dict("test" => 1))
     (val, state) = iterate(iterable, state)
@@ -288,6 +305,18 @@ test: 3
     @test iterate(iterable, state) === nothing
 end
 
+@testset "multi_doc_iteration_protocol" begin
+    fname = tempname() # cleanup=true, file will be deleted on process exit
+    open(fname, "w") do f
+        write(f, multidoc_contents)
+    end
+    iterable = YAML.load_all_file(fname)
+    @test Base.IteratorSize(YAML.YAMLDocIterator) == Base.SizeUnknown()
+    @test Base.IteratorEltype(YAML.YAMLDocIterator) == Base.HasEltype()
+    @test eltype(iterable) == Dict{Any, Any}
+    @test length(collect(iterable)) == 3
+end
+
 # test that an OrderedDict is written in the correct order
 using OrderedCollections, DataStructures
 @test strip(YAML.yaml(OrderedDict(:c => 3, :b => 2, :a => 1))) == join(["c: 3", "b: 2", "a: 1"], "\n")