| 111 | } |
| 112 | |
| 113 | func TestVisitFunc(t *testing.T) { |
| 114 | ast := &AST{ |
| 115 | Statements: []Statement{ |
| 116 | &SelectStatement{ |
| 117 | Columns: []Expression{ |
| 118 | &Identifier{Name: "id"}, |
| 119 | }, |
| 120 | TableName: "users", |
| 121 | }, |
| 122 | }, |
| 123 | } |
| 124 | |
| 125 | var count int |
| 126 | visitor := VisitFunc(func(n Node) (Visitor, error) { |
| 127 | if n != nil { |
| 128 | count++ |
| 129 | } |
| 130 | return nil, nil // Don't visit children |
| 131 | }) |
| 132 | |
| 133 | err := Walk(visitor, ast) |
| 134 | if err != nil { |
| 135 | t.Errorf("unexpected error: %v", err) |
| 136 | } |
| 137 | |
| 138 | if count != 1 { // Should only visit root node since we return nil |
| 139 | t.Errorf("expected count of 1, got %d", count) |
| 140 | } |
| 141 | } |