From 8217a8d86cbc4f5d0ea1c3f6e02dc271cbe3e73f Mon Sep 17 00:00:00 2001
From: Stephen Paul Weber <stephen.weber@shopify.com>
Date: Mon, 15 Oct 2018 10:22:09 -0400
Subject: [PATCH] Add test for the full array structure

---
 test/integration/parse_tree_visitor_test.rb | 78 ++++++++++++---------
 1 file changed, 46 insertions(+), 32 deletions(-)

diff --git a/test/integration/parse_tree_visitor_test.rb b/test/integration/parse_tree_visitor_test.rb
index 4ffef08..6ad6a2d 100644
--- a/test/integration/parse_tree_visitor_test.rb
+++ b/test/integration/parse_tree_visitor_test.rb
@@ -8,217 +8,228 @@ class ParseTreeVisitorTest < Minitest::Test
   def test_variable
     assert_equal(
       ["test"],
-      traversal(%({{ test }}))
+      visit(%({{ test }}))
     )
   end
 
   def test_varible_with_filter
     assert_equal(
       ["test", "infilter"],
-      traversal(%({{ test | split: infilter }}))
+      visit(%({{ test | split: infilter }}))
     )
   end
 
   def test_dynamic_variable
     assert_equal(
       ["test", "inlookup"],
-      traversal(%({{ test[inlookup] }}))
+      visit(%({{ test[inlookup] }}))
     )
   end
 
   def test_if_condition
     assert_equal(
       ["test"],
-      traversal(%({% if test %}{% endif %}))
+      visit(%({% if test %}{% endif %}))
     )
   end
 
   def test_complex_if_condition
     assert_equal(
       ["test"],
-      traversal(%({% if 1 == 1 and 2 == test %}{% endif %}))
+      visit(%({% if 1 == 1 and 2 == test %}{% endif %}))
     )
   end
 
   def test_if_body
     assert_equal(
       ["test"],
-      traversal(%({% if 1 == 1 %}{{ test }}{% endif %}))
+      visit(%({% if 1 == 1 %}{{ test }}{% endif %}))
     )
   end
 
   def test_unless_condition
     assert_equal(
       ["test"],
-      traversal(%({% unless test %}{% endunless %}))
+      visit(%({% unless test %}{% endunless %}))
     )
   end
 
   def test_complex_unless_condition
     assert_equal(
       ["test"],
-      traversal(%({% unless 1 == 1 and 2 == test %}{% endunless %}))
+      visit(%({% unless 1 == 1 and 2 == test %}{% endunless %}))
     )
   end
 
   def test_unless_body
     assert_equal(
       ["test"],
-      traversal(%({% unless 1 == 1 %}{{ test }}{% endunless %}))
+      visit(%({% unless 1 == 1 %}{{ test }}{% endunless %}))
     )
   end
 
   def test_elsif_condition
     assert_equal(
       ["test"],
-      traversal(%({% if 1 == 1 %}{% elsif test %}{% endif %}))
+      visit(%({% if 1 == 1 %}{% elsif test %}{% endif %}))
     )
   end
 
   def test_complex_elsif_condition
     assert_equal(
       ["test"],
-      traversal(%({% if 1 == 1 %}{% elsif 1 == 1 and 2 == test %}{% endif %}))
+      visit(%({% if 1 == 1 %}{% elsif 1 == 1 and 2 == test %}{% endif %}))
     )
   end
 
   def test_elsif_body
     assert_equal(
       ["test"],
-      traversal(%({% if 1 == 1 %}{% elsif 2 == 2 %}{{ test }}{% endif %}))
+      visit(%({% if 1 == 1 %}{% elsif 2 == 2 %}{{ test }}{% endif %}))
     )
   end
 
   def test_else_body
     assert_equal(
       ["test"],
-      traversal(%({% if 1 == 1 %}{% else %}{{ test }}{% endif %}))
+      visit(%({% if 1 == 1 %}{% else %}{{ test }}{% endif %}))
     )
   end
 
   def test_case_left
     assert_equal(
       ["test"],
-      traversal(%({% case test %}{% endcase %}))
+      visit(%({% case test %}{% endcase %}))
     )
   end
 
   def test_case_condition
     assert_equal(
       ["test"],
-      traversal(%({% case 1 %}{% when test %}{% endcase %}))
+      visit(%({% case 1 %}{% when test %}{% endcase %}))
     )
   end
 
   def test_case_when_body
     assert_equal(
       ["test"],
-      traversal(%({% case 1 %}{% when 2 %}{{ test }}{% endcase %}))
+      visit(%({% case 1 %}{% when 2 %}{{ test }}{% endcase %}))
     )
   end
 
   def test_case_else_body
     assert_equal(
       ["test"],
-      traversal(%({% case 1 %}{% else %}{{ test }}{% endcase %}))
+      visit(%({% case 1 %}{% else %}{{ test }}{% endcase %}))
     )
   end
 
   def test_for_in
     assert_equal(
       ["test"],
-      traversal(%({% for x in test %}{% endfor %}))
+      visit(%({% for x in test %}{% endfor %}))
     )
   end
 
   def test_for_limit
     assert_equal(
       ["test"],
-      traversal(%({% for x in (1..5) limit: test %}{% endfor %}))
+      visit(%({% for x in (1..5) limit: test %}{% endfor %}))
     )
   end
 
   def test_for_offset
     assert_equal(
       ["test"],
-      traversal(%({% for x in (1..5) offset: test %}{% endfor %}))
+      visit(%({% for x in (1..5) offset: test %}{% endfor %}))
     )
   end
 
   def test_for_body
     assert_equal(
       ["test"],
-      traversal(%({% for x in (1..5) %}{{ test }}{% endfor %}))
+      visit(%({% for x in (1..5) %}{{ test }}{% endfor %}))
     )
   end
 
   def test_tablerow_in
     assert_equal(
       ["test"],
-      traversal(%({% tablerow x in test %}{% endtablerow %}))
+      visit(%({% tablerow x in test %}{% endtablerow %}))
     )
   end
 
   def test_tablerow_limit
     assert_equal(
       ["test"],
-      traversal(%({% tablerow x in (1..5) limit: test %}{% endtablerow %}))
+      visit(%({% tablerow x in (1..5) limit: test %}{% endtablerow %}))
     )
   end
 
   def test_tablerow_offset
     assert_equal(
       ["test"],
-      traversal(%({% tablerow x in (1..5) offset: test %}{% endtablerow %}))
+      visit(%({% tablerow x in (1..5) offset: test %}{% endtablerow %}))
     )
   end
 
   def test_tablerow_body
     assert_equal(
       ["test"],
-      traversal(%({% tablerow x in (1..5) %}{{ test }}{% endtablerow %}))
+      visit(%({% tablerow x in (1..5) %}{{ test }}{% endtablerow %}))
     )
   end
 
   def test_cycle
     assert_equal(
       ["test"],
-      traversal(%({% cycle test %}))
+      visit(%({% cycle test %}))
     )
   end
 
   def test_assign
     assert_equal(
       ["test"],
-      traversal(%({% assign x = test %}))
+      visit(%({% assign x = test %}))
     )
   end
 
   def test_capture
     assert_equal(
       ["test"],
-      traversal(%({% capture x %}{{ test }}{% endcapture %}))
+      visit(%({% capture x %}{{ test }}{% endcapture %}))
     )
   end
 
   def test_include
     assert_equal(
       ["test"],
-      traversal(%({% include test %}))
+      visit(%({% include test %}))
     )
   end
 
   def test_include_with
     assert_equal(
       ["test"],
-      traversal(%({% include "hai" with test %}))
+      visit(%({% include "hai" with test %}))
     )
   end
 
   def test_include_for
     assert_equal(
       ["test"],
-      traversal(%({% include "hai" for test %}))
+      visit(%({% include "hai" for test %}))
+    )
+  end
+
+  def test_preserve_tree_structure
+    assert_equal(
+      [[nil, [
+        [nil, [[nil, [["other", []]]]]],
+        ["test", []],
+        ["xs", []]
+      ]]],
+      traversal(%({% for x in xs offset: test %}{{ other }}{% endfor %})).visit
     )
   end
 
@@ -228,6 +239,9 @@ class ParseTreeVisitorTest < Minitest::Test
     ParseTreeVisitor
       .for(Template.parse(template).root)
       .add_callback_for(VariableLookup, &:name)
-      .visit.flatten.compact
+  end
+
+  def visit(template)
+    traversal(template).visit.flatten.compact
   end
 end