Extracting ProjectStep from TinkerPop GraphTraversal

ProjectKeysExtractor.java

import static org.mockito.Mockito.mock;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.apache.tinkerpop.gremlin.process.traversal.Step;
import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.process.traversal.step.TraversalParent;
import org.apache.tinkerpop.gremlin.process.traversal.step.branch.LocalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.branch.OptionalStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.ProjectStep;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import lombok.Getter;

/**
 * This class sets up the necessary mocks and helper methods required to extract the keys (Strings)
 * provided to the {@code .project(...)} steps in a given {@link GraphTraversal}.
 * 
 * <p>
 * The method used to mock the TinkerPop classes is based on a <a href=
 * "https://github.com/apache/tinkerpop/blob/9fca27c0163d3ad56959adcf91aaa9215622adc6/gremlin-core/src/test/java/org/apache/tinkerpop/gremlin/process/traversal/dsl/graph/GraphTraversalSourceTest.java#L81-L127">method
 * published here</a>. It should be noted that this is very much a <strong>hack</strong>, and is
 * only used as a means to extract arguments passed to the {@code .project(...)} step. This strategy
 * should not be used to test the internals or validity of the query itself:
 * 
 * <blockquote><em>It is really important to note that all internal processing of the query is
 * ignored under this model and no TinkerPop code is executed so even invalid queries that are
 * incorrect will return the static data encoded to the mock</em> — TinkerPop GitHub</blockquote>
 * 
 * <p>
 * Some thoughts: Overall, this is put together scrappily. I didn't have much time to
 * look into the internals of how traversals are performed, but there might be a way to validate
 * that the actual Map being returned from the traversal is the one whose keys should be analysed.
 * In theory this should be possible, when you add .project() to a traversal, Java immediately knows
 * this know results in a Map, internally the ProjectStep itself might be encoded in some 'lastStep'
 * field.
 */
public class ProjectKeysExtractor {
  private final List<String> collectedProjectKeys = Lists.newArrayList();

  // These steps are considered 'parents' and may contain their own anonymous traversals
  private static final Map<String, Class<? extends TraversalParent>> PARENT_STEPS =
    ImmutableMap.<String, Class<? extends TraversalParent>>builder()
      .put("local", LocalStep.class)
      .put("optional", OptionalStep.class)
      .build();

  private GraphTraversal<?, ?> mockTraversal = mock(GraphTraversal.class, invocation -> {
    String methodName = invocation.getMethod().getName();
    if (methodName.equals("next")) {
      return this.collectedProjectKeys.stream().collect(Collectors.toList());
    }

    // Handle top-level project steps
    if (methodName.equals("project")) {
      Object[] args = invocation.getArguments();
      for (int i = 0; i < args.length; i += 1) {
        this.collectedProjectKeys.add(String.valueOf(args[i]));
      }
    }

    // Recursively find any project steps
    if (PARENT_STEPS.containsKey(methodName)) {
      Traversal<?, ?> localTraversal = (Traversal<?, ?>) invocation.getArguments()[0];
      this.collectedProjectKeys.addAll(extractProjectKeysFromTraversalSteps(localTraversal));
    }

    return invocation.getMock();
  });

  @Getter
  private GraphTraversalSource mockG = mock(GraphTraversalSource.class, invocation -> {
    final Class<?> returnType = invocation.getMethod()
      .getReturnType();
    if (returnType.isAssignableFrom(GraphTraversalSource.class)) {
      return invocation.getMock();
    }
    if (returnType.isAssignableFrom(GraphTraversal.class)) {
      return mockTraversal;
    }
    return invocation.callRealMethod();
  });

  /**
   * Handles recursively extracting project keys found in the provided traversal. Useful for
   * processing Anonymous Traversals.
   */
  private List<String> extractProjectKeysFromTraversalSteps(Traversal<?, ?> localTraversal) {
    List<String> projectKeys = Lists.newArrayList();

    for (final Step<?, ?> step : localTraversal.asAdmin().getSteps()) {
      if (PARENT_STEPS.containsValue(step.getClass())) {
        TraversalParent parentStep = (TraversalParent) step.clone();
        parentStep.getLocalChildren()
          .stream()
          .map(this::extractProjectKeysFromTraversalSteps)
          .forEach(projectKeys::addAll);
      }
      if (ProjectStep.class.isAssignableFrom(step.getClass())) {
        ProjectStep<?, ?> projectStep = (ProjectStep<?, ?>) step.clone();
        projectKeys.addAll(projectStep.getProjectKeys());
      }

    }
    return projectKeys;
  }

  @SuppressWarnings("unchecked")
  public List<String> extract(
      Function<GraphTraversalSource, GraphTraversal<?, ?>> traversalProvider) {
    this.collectedProjectKeys.clear();
    return (List<String>) traversalProvider.apply(this.mockG).next();
  }
}