gortiz commented on code in PR #14296: URL: https://github.com/apache/pinot/pull/14296#discussion_r1820964793
########## pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java: ########## @@ -0,0 +1,230 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.pinot.query.planner.logical; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.calcite.rel.core.JoinRelType; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.plannode.JoinNode; +import org.apache.pinot.query.planner.plannode.MailboxReceiveNode; +import org.apache.pinot.query.planner.plannode.MailboxSendNode; +import org.apache.pinot.query.planner.plannode.PlanNode; +import org.apache.pinot.query.planner.plannode.TableScanNode; +import org.testng.annotations.AfterMethod; + + +/** + * A base test class that can be used to write tests for stages using a fluent DSL. + * + * This class provides two features: + * <ul> + * <li>Builders that can be used to create nodes in a fluent way.</li> + * <li>Access to the stages that were created during the test with {@link #stage(int)}.</li> + * </ul> + * + * It is expected that each test method will call {@link #when(SimpleChildBuilder)} to create a new plan, which will + * populate the list of stages. After that, the test can look for the stages with the {@link #stage(int)} method and + * assert the expected behavior. + */ +public class StagesTestBase { + private final HashMap<Integer, MailboxSendNode> _stageRoots = new HashMap<>(); + + /** + * Clears the list of stages. + * + * This method is automatically called by the test framework, ensuring each test starts with a clean slate. + * This method can also be called in middle of the test, but that is not recommended given it usually means that the + * test is getting too complex and difficult to read and/or get insights from it in case of failure. + */ + @AfterMethod Review Comment: Good catch! I wrote it to be `@BeforeMethod` but I made the mistake of using `@BeforeTest` (I hate TestNG API so much! :D). When I fixed it I thought `@AfterMethod` is in theory better given it is a cleanup that should be executed after each test. I can rename it or change the hook -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org
