eu.drus.jpa.unit.neo4j.CleanupStrategyProvider Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jpa-unit-neo4j Show documentation
JUnit extension for simple testing of JPA entities and components
The newest version!
package eu.drus.jpa.unit.neo4j;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;

import org.jgrapht.DirectedGraph;
import org.jgrapht.Graph;
import org.jgrapht.Graphs;
import org.jgrapht.graph.ClassBasedEdgeFactory;
import org.jgrapht.graph.DefaultDirectedGraph;

import eu.drus.jpa.unit.api.CleanupStrategy.StrategyProvider;
import eu.drus.jpa.unit.neo4j.dataset.DatabaseReader;
import eu.drus.jpa.unit.neo4j.dataset.Edge;
import eu.drus.jpa.unit.neo4j.dataset.GraphElementFactory;
import eu.drus.jpa.unit.neo4j.dataset.Node;
import eu.drus.jpa.unit.neo4j.operation.Neo4JOperations;
import eu.drus.jpa.unit.spi.CleanupStrategyExecutor;
import eu.drus.jpa.unit.spi.DbFeatureException;

public class CleanupStrategyProvider implements StrategyProvider>> {

    private static final String UNABLE_TO_CLEAN_DATABASE = "Unable to clean database.";
    private DatabaseReader dbReader;

    public CleanupStrategyProvider(final GraphElementFactory factory) {
        dbReader = new DatabaseReader(factory);
    }

    @Override
    public CleanupStrategyExecutor> strictStrategy() {
        return (final Connection connection, final List> initialGraphs, final String... nodeTypesToRetain) -> {

            try {
                Neo4JOperations.DELETE_ALL.execute(connection, computeGraphToBeDeleted(dbReader.readGraph(connection), nodeTypesToRetain));
                connection.commit();
            } catch (final SQLException e) {
                throw new DbFeatureException(UNABLE_TO_CLEAN_DATABASE, e);
            }
        };
    }

    @Override
    public CleanupStrategyExecutor> usedTablesOnlyStrategy() {
        return (final Connection connection, final List> initialGraphs, final String... nodeTypesToRetain) -> {
            if (initialGraphs.isEmpty()) {
                return;
            }

            try {
                for (final Graph graph : initialGraphs) {
                    Neo4JOperations.DELETE_ALL.execute(connection, computeGraphToBeDeleted(graph, nodeTypesToRetain));
                }

                connection.commit();
            } catch (final SQLException e) {
                throw new DbFeatureException(UNABLE_TO_CLEAN_DATABASE, e);
            }
        };
    }

    @Override
    public CleanupStrategyExecutor> usedRowsOnlyStrategy() {
        return (final Connection connection, final List> initialGraphs, final String... nodeTypesToRetain) -> {
            if (initialGraphs.isEmpty()) {
                return;
            }

            try {
                for (final Graph graph : initialGraphs) {
                    Neo4JOperations.DELETE.execute(connection, computeGraphToBeDeleted(graph, nodeTypesToRetain));
                }

                connection.commit();
            } catch (final SQLException e) {
                throw new DbFeatureException(UNABLE_TO_CLEAN_DATABASE, e);
            }
        };
    }

    private Graph computeGraphToBeDeleted(final Graph graph, final String... nodeTypesToRetain) {
        final DirectedGraph toDelete = new DefaultDirectedGraph<>(new ClassBasedEdgeFactory<>(Edge.class));

        // copy graph to a destination, which we are going to modify
        Graphs.addGraph(toDelete, graph);

        // remove the nodes, we have to retain from the graph
        Graphs.removeVerticesAndPreserveConnectivity(toDelete, v -> shouldRetainNode(v, nodeTypesToRetain));

        return toDelete;
    }

    private boolean shouldRetainNode(final Node node, final String... nodeTypesToRetain) {
        for (final String nodeToExclude : nodeTypesToRetain) {
            if (node.getLabels().contains(nodeToExclude)) {
                return true;
            }
        }
        return false;
    }
}