Unit Testing Apache TinkerPop Transactions: From TinkerGraph to Amazon Neptune

<dependency>
	<groupId>org.apache.tinkerpop</groupId>
	<artifactId>tinkergraph-gremlin</artifactId>
	<version>3.7.0</version>
	<scope>test</scope>
</dependency>

public interface AirportService {
    public boolean addAirport(Map<String, Object> airportData);
    public boolean addRoute(String fromAirport, String toAirport, int distance);
    public Map<String, Object> getAirportData(String airportCode);
    public int getRouteDistance(String fromAirportCode, String toAirportCode);
    public boolean hasRoute(String fromAirport, String toAirport);
    public boolean removeAirport(String airportCode);
    public boolean removeRoute(String fromAirportCode, String toAirportCode);
}

public class NorthAmericanAirportService implements AirportService {
    private GraphTraversalSource g;

    public NorthAmericanAirportService(GraphTraversalSource g) {
        this.g = g;
    }
    
    /**
     * Adds a route between two airports.
     *
     * @param fromAirportCode   The airport code of airport where the route begins.
     * @param toAirportCode     The airport code of airport where the route ends.
     * @param distance          The distance between the two airports.
     * @return                  True if the route was added; false otherwise.
     */
    public boolean addRoute(String fromAirportCode, String toAirportCode, int distance) {
        Transaction tx = g.tx();
        GraphTraversalSource gtx = tx.begin(); // Explicitly starting the transaction.

        // This try-catch-rollback approach is recommended with TinkerPop transactions.
        try {
            final Vertex fromV = gtx.V().has("code", fromAirportCode).next();
            final Vertex toV = gtx.V().has("code", toAirportCode).next();
            gtx.addE("route").from(fromV).to(toV).next();
            tx.commit();

            return true;
        } catch (Exception e) {
            tx.rollback();
            return false;
        }
    }

public class AirportServiceTest {

    // In this example, "STAGING_ENV" is used to determine whether to test against TinkerGraph or Amazon Neptune.
    private static boolean STAGING_ENV = (null != System.getProperty("STAGING_ENV"));
    private static Cluster cluster;

    private GraphTraversalSource g;

    @BeforeClass
    public static void setUpServerCluster() {
        if (STAGING_ENV) {
            cluster = Cluster.build().addContactPoint("your-neptune-cluster").enableSsl(true).create();
        }
    }
    
    @Before
    public void setUpGraph() {
        if (STAGING_ENV) { // In this example, STAGING_ENV is a system property used to determine which database to use.
            g = traversal().withRemote(DriverRemoteConnection.using(cluster));
            g.V().drop().iterate();
            // Currently, Neptune only accepts URLs reachable within its VPC as an input to the io step.
            g.io("your-presigned-s3-url").read().iterate();
        } else {
            // Create a default, empty instance of the transactional TinkerGraph.
            g = TinkerTransactionGraph.open().traversal();

            g.io("your-local-data.xml").read().iterate(); // This is how you insert your GraphML test data.
            g.tx().commit(); // By default, transactions are automatically opened, so commit the changes from io().
        }
    }

    @Test
    public void testAddRouteWithIncorrectAirportCode() {
        final NorthAmericanAirportService service = new NorthAmericanAirportService(g);

        // Add route with airport code that doesn't exist.
        final boolean wasAdded = service.addRoute("INCORRECT", "AUS", 500);

        assertFalse(wasAdded);
        // Check to see if any routes exist between those two airports.
        assertEquals(0L,
                     g.E().where(inV().has("code", "INCORRECT"))
                          .where(outV().has("code", "AUS")).count().next().longValue());
    }
	@Test
    public void testAddRouteWithValidAirportCodes() {
        final NorthAmericanAirportService service = new NorthAmericanAirportService(g);

        final boolean wasAdded = service.addRoute("PBI", "ORD", 500);

        assertTrue(wasAdded);
        assertEquals(1L,
                     g.E().where(inV().has("code", "PBI"))
                          .where(outV().has("code", "ORD")).count().next().longValue());
    }

/**
     * Removes incoming routes to an airport.
     *
     * @param airportCode   The airport code of the airport to remove incoming routes.
     */
    public void stopIncomingTraffic(String airportCode) {
        Transaction tx = g.tx();

        try {
            g.V().has("code", airportCode).inE().drop().iterate();
            tx.commit();
        } catch (Exception e) {
            tx.rollback();
        }
    }

@Test
    public void testStoppingTrafficToAus() {
        final NorthAmericanAirportService service = new NorthAmericanAirportService(g);
        final String airport = "AUS";

        service.stopIncomingTraffic(airport);

        // Check that there are no outgoing routes into that airport.
        assertEquals(0, g.V().out().has("code", airport).toList().size());
    }