/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.prestosql.sql.planner.iterative.rule.test;

import com.google.common.collect.ImmutableList;
import io.prestosql.matching.Captures;
import io.prestosql.matching.Pattern;
import io.prestosql.sql.planner.assertions.PlanMatchPattern;
import io.prestosql.sql.planner.iterative.Rule;
import io.prestosql.sql.planner.plan.Assignments;
import io.prestosql.sql.planner.plan.PlanNode;
import org.testng.annotations.Test;

import static io.prestosql.sql.planner.assertions.PlanMatchPattern.values;
import static io.prestosql.sql.planner.iterative.rule.test.PlanBuilder.expression;
import static io.prestosql.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

public class TestRuleTester
{
    @Test
    public void testReportWrongMatch()
    {
        try (RuleTester tester = defaultRuleTester()) {
            RuleAssert ruleAssert = tester.assertThat(new DummyReplaceNodeRule())
                    .on(p ->
                            p.project(
                                    Assignments.of(p.symbol("y"), expression("x")),
                                    p.values(
                                            ImmutableList.of(p.symbol("x")),
                                            ImmutableList.of(ImmutableList.of(expression("1"))))));

            PlanMatchPattern expected = values(ImmutableList.of("different"), ImmutableList.of());
            assertThatThrownBy(() -> ruleAssert.matches(expected))
                    .isInstanceOf(AssertionError.class)
                    .hasMessageMatching("(?s)Plan does not match, expected .* but found .*");
        }
    }

    private static class DummyReplaceNodeRule
            implements Rule<PlanNode>
    {
        @Override
        public Pattern<PlanNode> getPattern()
        {
            return Pattern.typeOf(PlanNode.class);
        }

        @Override
        public Result apply(PlanNode node, Captures captures, Context context)
        {
            return Result.ofPlanNode(node.replaceChildren(node.getSources()));
        }
    }
}

@@ -844,16 +844,127 @@
 e.Rule;%0A
+import io.prestosql.sql.planner.iterative.Rule.Context;%0Aimport io.prestosql.sql.planner.iterative.Rule.Result;%0A
 import i
@@ -1336,16 +1336,64 @@
 Tester;%0A
+import static java.util.Objects.requireNonNull;%0A
 import s
@@ -1650,33 +1650,267 @@
 hat(
-new DummyReplaceNodeRule(
+%0A                    rule(%0A                            %22testReportWrongMatch rule%22,%0A                            Pattern.typeOf(PlanNode.class),%0A                            (node, captures, context) -%3E Result.ofPlanNode(node.replaceChildren(node.getSources()))
 ))%0A 
@@ -2640,47 +2640,106 @@
 tic 
-class DummyReplaceNodeRule
+%3CT%3E Rule%3CT%3E rule(String name, Pattern%3CT%3E pattern, RuleApplyImplementation%3CT%3E apply)
 %0A    
+%7B%0A
         
 impl
@@ -2738,34 +2738,179 @@
     
-implements Rule%3CPlanNode%3E%0A
+requireNonNull(name, %22name is null%22);%0A        requireNonNull(pattern, %22pattern is null%22);%0A        requireNonNull(apply, %22apply is null%22);%0A        return new Rule%3CT%3E()%0A    
     
@@ -2911,32 +2911,36 @@
       %7B%0A        
+    
 @Override%0A      
@@ -2945,31 +2945,145 @@
     
+    
 public 
-Pattern%3CPlanNode
+String toString()%0A            %7B%0A                return name;%0A            %7D%0A%0A            @Override%0A            public Pattern%3CT
 %3E ge
@@ -3093,34 +3093,42 @@
 ttern()%0A        
+    
 %7B%0A
+    
             retu
@@ -3134,51 +3134,36 @@
 urn 
-P
+p
 attern
-.typeOf(PlanNode.class);%0A        %7D%0A%0A
+;%0A            %7D%0A%0A    
     
@@ -3180,24 +3180,28 @@
 ide%0A        
+    
 public Resul
@@ -3208,24 +3208,17 @@
 t apply(
-PlanNode
+T
  node, C
@@ -3252,34 +3252,42 @@
 ontext)%0A        
-%7B%0A
+    %7B%0A    
             retu
@@ -3293,77 +3293,223 @@
 urn 
-Result.ofPlanNode(node.replaceChildren(node.getSources()));%0A        %7D
+apply.apply(node, captures, context);%0A            %7D%0A        %7D;%0A    %7D%0A%0A    @FunctionalInterface%0A    private interface RuleApplyImplementation%3CT%3E%0A    %7B%0A        Result apply(T node, Captures captures, Context context);
 %0A   

Move helper rule implementation to the test method