<commit_before>/*
 * Copyright 2016-17 Crown Copyright
 *
 * 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 uk.gov.gchq.gaffer.sparkaccumulo.operation.handler.scalardd;

import com.google.common.collect.Sets;
import org.apache.accumulo.core.client.AccumuloException;
import org.apache.accumulo.core.client.AccumuloSecurityException;
import org.apache.accumulo.core.client.TableNotFoundException;
import org.apache.accumulo.core.client.admin.CompactionConfig;
import org.apache.accumulo.minicluster.MiniAccumuloCluster;
import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.hadoop.conf.Configuration;
import org.apache.spark.rdd.RDD;
import org.apache.spark.sql.SparkSession;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import uk.gov.gchq.gaffer.accumulostore.AccumuloProperties;
import uk.gov.gchq.gaffer.accumulostore.AccumuloStore;
import uk.gov.gchq.gaffer.commonutil.CommonConstants;
import uk.gov.gchq.gaffer.commonutil.CommonTestConstants;
import uk.gov.gchq.gaffer.commonutil.TestGroups;
import uk.gov.gchq.gaffer.commonutil.TestPropertyNames;
import uk.gov.gchq.gaffer.data.element.Edge;
import uk.gov.gchq.gaffer.data.element.Element;
import uk.gov.gchq.gaffer.data.element.Entity;
import uk.gov.gchq.gaffer.data.element.IdentifierType;
import uk.gov.gchq.gaffer.data.element.function.ElementTransformer;
import uk.gov.gchq.gaffer.data.elementdefinition.view.View;
import uk.gov.gchq.gaffer.data.elementdefinition.view.ViewElementDefinition;
import uk.gov.gchq.gaffer.graph.Graph;
import uk.gov.gchq.gaffer.graph.GraphConfig;
import uk.gov.gchq.gaffer.operation.OperationException;
import uk.gov.gchq.gaffer.operation.impl.add.AddElements;
import uk.gov.gchq.gaffer.serialisation.implementation.StringSerialiser;
import uk.gov.gchq.gaffer.spark.SparkConstants;
import uk.gov.gchq.gaffer.spark.operation.scalardd.GetRDDOfAllElements;
import uk.gov.gchq.gaffer.sparkaccumulo.operation.handler.AbstractGetRDDHandler;
import uk.gov.gchq.gaffer.sparkaccumulo.operation.handler.MiniAccumuloClusterProvider;
import uk.gov.gchq.gaffer.store.schema.Schema;
import uk.gov.gchq.gaffer.store.schema.SchemaEntityDefinition;
import uk.gov.gchq.gaffer.store.schema.TypeDefinition;
import uk.gov.gchq.gaffer.user.User;
import uk.gov.gchq.koryphe.impl.function.Concat;

import java.io.DataOutputStream;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

public class GetRDDOfAllElementsHandlerTest {
    private final User USER = new User();
    private final User USER_WITH_PUBLIC = new User("user1", Sets.newHashSet("public"));
    private final User USER_WITH_PUBLIC_AND_PRIVATE = new User("user2", Sets.newHashSet("public", "private"));
    private final String GRAPH_ID = "graphId";
    private Entity entityRetainedAfterValidation;

    @Test
    public void testGetAllElementsInRDD() throws OperationException, IOException, InterruptedException,
            AccumuloSecurityException, AccumuloException, TableNotFoundException {
        testGetAllElementsInRDD(getGraphForMockAccumulo(), getOperation());
        testGetAllElementsInRDD(getGraphForDirectRDD("testGetAllElementsInRDD"), getOperationWithDirectRDDOption());
    }

    @Test
    public void testGetAllElementsInRDDWithView() throws OperationException, IOException, InterruptedException,
            AccumuloSecurityException, AccumuloException, TableNotFoundException {
        testGetAllElementsInRDDWithView(getGraphForMockAccumulo(), getOperation());
        testGetAllElementsInRDDWithView(getGraphForDirectRDD("testGetAllElementsInRDDWithView"), getOperationWithDirectRDDOption());
    }

    @Test
    public void testGetAllElementsInRDDWithVisibilityFilteringApplied() throws OperationException, IOException,
            InterruptedException, AccumuloSecurityException, AccumuloException, TableNotFoundException {
        testGetAllElementsInRDDWithVisibilityFilteringApplied(getGraphForMockAccumuloWithVisibility(), getOperation());
        testGetAllElementsInRDDWithVisibilityFilteringApplied(
                getGraphForDirectRDDWithVisibility("testGetAllElementsInRDDWithVisibilityFilteringApplied"),
                getOperationWithDirectRDDOption());
    }

    @Test
    public void testGetAllElementsInRDDWithValidationApplied() throws InterruptedException, IOException,
            OperationException, AccumuloSecurityException, TableNotFoundException, AccumuloException {
        testGetAllElementsInRDDWithValidationApplied(getGraphForMockAccumuloForValidationChecking(), getOperation());
        testGetAllElementsInRDDWithValidationApplied(
                getGraphForDirectRDDForValidationChecking("testGetAllElementsInRDDWithValidationApplied"),
                getOperationWithDirectRDDOption());
    }

    private void testGetAllElementsInRDD(final Graph graph, final GetRDDOfAllElements getRDD) throws OperationException,
            IOException, InterruptedException, AccumuloSecurityException, AccumuloException {
        final Set<Element> expectedElements = new HashSet<>(getElements());
        final RDD<Element> rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }
        final Set<Element> results = new HashSet<>();
        final Element[] returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        getRDD.getSparkSession().stop();
    }

    private void testGetAllElementsInRDDWithView(final Graph graph, final GetRDDOfAllElements getRDD) throws OperationException,
            IOException, InterruptedException, AccumuloSecurityException, AccumuloException {
        final Set<Element> expectedElements = new HashSet<>();
        getElements().stream()
                .filter(e -> e.getGroup().equals(TestGroups.EDGE))
                .map(e -> (Edge) e)
                .map(e -> {
                    e.putProperty("newProperty", e.getSource().toString() + "," + e.getProperty(TestPropertyNames.COUNT));
                    return e;
                })
                .forEach(expectedElements::add);
        getRDD.setView(new View.Builder()
                .edge(TestGroups.EDGE, new ViewElementDefinition.Builder()
                        .transientProperty("newProperty", String.class)
                        .transformer(new ElementTransformer.Builder()
                                .select(IdentifierType.SOURCE.name(), TestPropertyNames.COUNT)
                                .execute(new Concat())
                                .project("newProperty")
                                .build())
                        .build())
                .build());
        final RDD<Element> rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }
        final Set<Element> results = new HashSet<>();
        final Element[] returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        getRDD.getSparkSession().stop();
    }

    private void testGetAllElementsInRDDWithVisibilityFilteringApplied(final Graph graph,
                                                                       final GetRDDOfAllElements getRDD)
            throws OperationException, IOException, InterruptedException, AccumuloSecurityException, AccumuloException {
        final Set<Element> expectedElements = new HashSet<>();

        // Test with user with public visibility
        getElementsWithVisibilities()
                .stream()
                .filter(e -> e.getProperty(TestPropertyNames.VISIBILITY).equals("public"))
                .forEach(expectedElements::add);
        RDD<Element> rdd = graph.execute(getRDD, USER_WITH_PUBLIC);
        if (rdd == null) {
            fail("No RDD returned");
        }
        final Set<Element> results = new HashSet<>();
        Element[] returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        // Test with user with public and private visibility
        getElementsWithVisibilities().forEach(expectedElements::add);
        rdd = graph.execute(getRDD, USER_WITH_PUBLIC_AND_PRIVATE);
        if (rdd == null) {
            fail("No RDD returned");
        }
        results.clear();
        returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        // Test with user with no visibilities
        rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }
        results.clear();
        returnedElements = (Element[]) rdd.collect();
        assertEquals(0, returnedElements.length);

        getRDD.getSparkSession().stop();
    }

    private void testGetAllElementsInRDDWithValidationApplied(final Graph graph, final GetRDDOfAllElements getRDD)
            throws InterruptedException, IOException, OperationException {
        // Sleep for 1 second to give chance for Entity A to age off
        Thread.sleep(1000L);

        final RDD<Element> rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }

        // Should get Entity B but not Entity A
        final Element[] returnedElements = (Element[]) rdd.collect();
        assertEquals(1, returnedElements.length);

        assertEquals(entityRetainedAfterValidation, returnedElements[0]);

        getRDD.getSparkSession().stop();
    }

    private Graph getGraphForMockAccumulo() throws OperationException {
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(GRAPH_ID)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elements.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getClass().getResourceAsStream("/store.properties"))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElements()).build(), user);
        return graph;
    }

    private Graph getGraphForMockAccumuloWithVisibility() throws OperationException {
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(GRAPH_ID)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsWithVisibility.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getClass().getResourceAsStream("/store.properties"))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsWithVisibilities()).build(), user);
        return graph;
    }

    private Graph getGraphForMockAccumuloForValidationChecking() throws OperationException {
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(GRAPH_ID)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/typesForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getClass().getResourceAsStream("/store.properties"))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsForValidationChecking()).build(), user);
        return graph;
    }

    private Graph getGraphForDirectRDD(final String tableName) throws InterruptedException, AccumuloException, AccumuloSecurityException,
            IOException, OperationException, TableNotFoundException {
        final MiniAccumuloCluster cluster = MiniAccumuloClusterProvider.getMiniAccumuloCluster();
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(tableName)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elements.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getAccumuloProperties(cluster))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElements()).build(), user);
        cluster.getConnector("root", MiniAccumuloClusterProvider.PASSWORD)
                .tableOperations()
                .compact(tableName, new CompactionConfig());
        Thread.sleep(1000L);
        return graph;
    }

    private Graph getGraphForDirectRDDWithVisibility(final String tableName) throws InterruptedException, AccumuloException,
            AccumuloSecurityException, IOException, OperationException, TableNotFoundException {
        final MiniAccumuloCluster cluster = MiniAccumuloClusterProvider.getMiniAccumuloCluster();
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(tableName)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsWithVisibility.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getAccumuloProperties(cluster))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsWithVisibilities()).build(), user);
        cluster.getConnector("root", MiniAccumuloClusterProvider.PASSWORD)
                .tableOperations()
                .compact(tableName, new CompactionConfig());
        Thread.sleep(1000L);
        return graph;
    }

    private Graph getGraphForDirectRDDForValidationChecking(final String tableName) throws InterruptedException,
            AccumuloException, AccumuloSecurityException, IOException, OperationException, TableNotFoundException {
        final MiniAccumuloCluster cluster = MiniAccumuloClusterProvider.getMiniAccumuloCluster();
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(tableName)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/typesForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getAccumuloProperties(cluster))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsForValidationChecking()).build(), user);
        cluster.getConnector("root", MiniAccumuloClusterProvider.PASSWORD)
                .tableOperations()
                .compact(tableName, new CompactionConfig());
        Thread.sleep(1000L);
        return graph;
    }

    private List<Element> getElements() {
        final List<Element> elements = new ArrayList<>();
        for (int i = 0; i < 1; i++) {
            final Entity entity = new Entity.Builder()
                    .group(TestGroups.ENTITY)
                    .vertex("" + i)
                    .build();

            final Edge edge1 = new Edge.Builder()
                    .group(TestGroups.EDGE)
                    .source("" + i)
                    .dest("B")
                    .directed(false)
                    .property(TestPropertyNames.COUNT, 2)
                    .build();

            final Edge edge2 = new Edge.Builder()
                    .group(TestGroups.EDGE)
                    .source("" + i)
                    .dest("C")
                    .directed(false)
                    .property(TestPropertyNames.COUNT, 4)
                    .build();

            elements.add(edge1);
            elements.add(edge2);
            elements.add(entity);
        }
        return elements;
    }

    private List<Element> getElementsWithVisibilities() {
        return getElements().stream()
                .map(e -> {
                    if (e.getGroup().equals(TestGroups.ENTITY)) {
                        e.putProperty(TestPropertyNames.VISIBILITY, "public");
                    } else if (((Edge) e).getDestination().equals("B")) {
                        e.putProperty(TestPropertyNames.VISIBILITY, "private");
                    } else {
                        e.putProperty(TestPropertyNames.VISIBILITY, "public");
                    }
                    return e;
                })
                .collect(Collectors.toList());
    }

    private List<Element> getElementsForValidationChecking() {
        final List<Element> elements = new ArrayList<>();
        final Entity entity1 = new Entity.Builder()
                .group(TestGroups.ENTITY)
                .vertex("A")
                .property("timestamp", System.currentTimeMillis())
                .build();
        final Entity entity2 = new Entity.Builder()
                .group(TestGroups.ENTITY)
                .vertex("B")
                .property("timestamp", System.currentTimeMillis() + 1000000L)
                .build();
        entityRetainedAfterValidation = entity2;
        elements.add(entity1);
        elements.add(entity2);
        return elements;
    }

    private GetRDDOfAllElements getOperation() throws IOException {
        // Create Hadoop configuration and serialise to a string
        final Configuration configuration = new Configuration();
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        configuration.write(new DataOutputStream(baos));
        final String configurationString = new String(baos.toByteArray(), CommonConstants.UTF_8);

        // Check get correct elements
        final GetRDDOfAllElements rddQuery = new GetRDDOfAllElements.Builder()
                .sparkSession(getSparkSession())
                .build();
        rddQuery.addOption(AbstractGetRDDHandler.HADOOP_CONFIGURATION_KEY, configurationString);
        return rddQuery;
    }

    private GetRDDOfAllElements getOperationWithDirectRDDOption() throws IOException {
        final GetRDDOfAllElements op = getOperation();
        op.addOption(AbstractGetRDDHandler.USE_RFILE_READER_RDD, "true");
        return op;
    }

    private AccumuloProperties getAccumuloProperties(final MiniAccumuloCluster cluster) {
        final AccumuloProperties properties = new AccumuloProperties();
        properties.setStoreClass(AccumuloStore.class);
        properties.setInstance(cluster.getInstanceName());
        properties.setZookeepers(cluster.getZooKeepers());
        properties.setUser(MiniAccumuloClusterProvider.USER);
        properties.setPassword(<PASSWORD>);
        properties.setOperationDeclarationPaths("sparkAccumuloOperationsDeclarations.json");
        return properties;
    }

    private SparkSession getSparkSession() {
        final SparkSession sparkSession = SparkSession.builder()
                .master("local")
                .appName("testGetAllElementsInRDD-" + System.currentTimeMillis())
                .config(SparkConstants.SERIALIZER, SparkConstants.DEFAULT_SERIALIZER)
                .config(SparkConstants.KRYO_REGISTRATOR, SparkConstants.DEFAULT_KRYO_REGISTRATOR)
                .config(SparkConstants.DRIVER_ALLOW_MULTIPLE_CONTEXTS, true)
                .getOrCreate();
        return sparkSession;
    }
}<commit_msg>gh-1038 Fix issue in GetRDDOfAllElementsHandlerTest<commit_after>/*
 * Copyright 2016-17 Crown Copyright
 *
 * 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 uk.gov.gchq.gaffer.sparkaccumulo.operation.handler.scalardd;

import com.google.common.collect.Sets;
import org.apache.accumulo.core.client.AccumuloException;
import org.apache.accumulo.core.client.AccumuloSecurityException;
import org.apache.accumulo.core.client.TableNotFoundException;
import org.apache.accumulo.core.client.admin.CompactionConfig;
import org.apache.accumulo.minicluster.MiniAccumuloCluster;
import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.hadoop.conf.Configuration;
import org.apache.spark.rdd.RDD;
import org.apache.spark.sql.SparkSession;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import uk.gov.gchq.gaffer.accumulostore.AccumuloProperties;
import uk.gov.gchq.gaffer.accumulostore.AccumuloStore;
import uk.gov.gchq.gaffer.commonutil.CommonConstants;
import uk.gov.gchq.gaffer.commonutil.CommonTestConstants;
import uk.gov.gchq.gaffer.commonutil.TestGroups;
import uk.gov.gchq.gaffer.commonutil.TestPropertyNames;
import uk.gov.gchq.gaffer.data.element.Edge;
import uk.gov.gchq.gaffer.data.element.Element;
import uk.gov.gchq.gaffer.data.element.Entity;
import uk.gov.gchq.gaffer.data.element.IdentifierType;
import uk.gov.gchq.gaffer.data.element.function.ElementTransformer;
import uk.gov.gchq.gaffer.data.elementdefinition.view.View;
import uk.gov.gchq.gaffer.data.elementdefinition.view.ViewElementDefinition;
import uk.gov.gchq.gaffer.graph.Graph;
import uk.gov.gchq.gaffer.graph.GraphConfig;
import uk.gov.gchq.gaffer.operation.OperationException;
import uk.gov.gchq.gaffer.operation.impl.add.AddElements;
import uk.gov.gchq.gaffer.serialisation.implementation.StringSerialiser;
import uk.gov.gchq.gaffer.spark.SparkConstants;
import uk.gov.gchq.gaffer.spark.operation.scalardd.GetRDDOfAllElements;
import uk.gov.gchq.gaffer.sparkaccumulo.operation.handler.AbstractGetRDDHandler;
import uk.gov.gchq.gaffer.sparkaccumulo.operation.handler.MiniAccumuloClusterProvider;
import uk.gov.gchq.gaffer.store.schema.Schema;
import uk.gov.gchq.gaffer.store.schema.SchemaEntityDefinition;
import uk.gov.gchq.gaffer.store.schema.TypeDefinition;
import uk.gov.gchq.gaffer.user.User;
import uk.gov.gchq.koryphe.impl.function.Concat;

import java.io.DataOutputStream;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

public class GetRDDOfAllElementsHandlerTest {
    private final User USER = new User();
    private final User USER_WITH_PUBLIC = new User("user1", Sets.newHashSet("public"));
    private final User USER_WITH_PUBLIC_AND_PRIVATE = new User("user2", Sets.newHashSet("public", "private"));
    private final String GRAPH_ID = "graphId";
    private Entity entityRetainedAfterValidation;

    @Test
    public void testGetAllElementsInRDD() throws OperationException, IOException, InterruptedException,
            AccumuloSecurityException, AccumuloException, TableNotFoundException {
        testGetAllElementsInRDD(getGraphForMockAccumulo(), getOperation("testGetAllElementsInRDD2"));
        testGetAllElementsInRDD(getGraphForDirectRDD("testGetAllElementsInRDD"),
                getOperationWithDirectRDDOption("testGetAllElementsInRDD2"));
    }

    @Test
    public void testGetAllElementsInRDDWithView() throws OperationException, IOException, InterruptedException,
            AccumuloSecurityException, AccumuloException, TableNotFoundException {
        testGetAllElementsInRDDWithView(getGraphForMockAccumulo(), getOperation("testGetAllElementsInRDDWithView1"));
        testGetAllElementsInRDDWithView(getGraphForDirectRDD("testGetAllElementsInRDDWithView"),
                getOperationWithDirectRDDOption("testGetAllElementsInRDDWithView2"));
    }

    @Test
    public void testGetAllElementsInRDDWithVisibilityFilteringApplied() throws OperationException, IOException,
            InterruptedException, AccumuloSecurityException, AccumuloException, TableNotFoundException {
        testGetAllElementsInRDDWithVisibilityFilteringApplied(getGraphForMockAccumuloWithVisibility(),
                getOperation("testGetAllElementsInRDDWithVisibilityFilteringApplied1"));
        testGetAllElementsInRDDWithVisibilityFilteringApplied(
                getGraphForDirectRDDWithVisibility("testGetAllElementsInRDDWithVisibilityFilteringApplied"),
                getOperationWithDirectRDDOption("testGetAllElementsInRDDWithVisibilityFilteringApplied2"));
    }

    @Test
    public void testGetAllElementsInRDDWithValidationApplied() throws InterruptedException, IOException,
            OperationException, AccumuloSecurityException, TableNotFoundException, AccumuloException {
        testGetAllElementsInRDDWithValidationApplied(getGraphForMockAccumuloForValidationChecking(),
                getOperation("testGetAllElementsInRDDWithVisibilityFilteringApplied1"));
        testGetAllElementsInRDDWithValidationApplied(
                getGraphForDirectRDDForValidationChecking("testGetAllElementsInRDDWithValidationApplied"),
                getOperationWithDirectRDDOption("testGetAllElementsInRDDWithVisibilityFilteringApplied2"));
    }

    private void testGetAllElementsInRDD(final Graph graph, final GetRDDOfAllElements getRDD) throws OperationException,
            IOException, InterruptedException, AccumuloSecurityException, AccumuloException {
        final Set<Element> expectedElements = new HashSet<>(getElements());
        final RDD<Element> rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }
        final Set<Element> results = new HashSet<>();
        final Element[] returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        getRDD.getSparkSession().stop();
    }

    private void testGetAllElementsInRDDWithView(final Graph graph, final GetRDDOfAllElements getRDD) throws OperationException,
            IOException, InterruptedException, AccumuloSecurityException, AccumuloException {
        final Set<Element> expectedElements = new HashSet<>();
        getElements().stream()
                .filter(e -> e.getGroup().equals(TestGroups.EDGE))
                .map(e -> (Edge) e)
                .map(e -> {
                    e.putProperty("newProperty", e.getSource().toString() + "," + e.getProperty(TestPropertyNames.COUNT));
                    return e;
                })
                .forEach(expectedElements::add);
        getRDD.setView(new View.Builder()
                .edge(TestGroups.EDGE, new ViewElementDefinition.Builder()
                        .transientProperty("newProperty", String.class)
                        .transformer(new ElementTransformer.Builder()
                                .select(IdentifierType.SOURCE.name(), TestPropertyNames.COUNT)
                                .execute(new Concat())
                                .project("newProperty")
                                .build())
                        .build())
                .build());
        final RDD<Element> rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }
        final Set<Element> results = new HashSet<>();
        final Element[] returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        getRDD.getSparkSession().stop();
    }

    private void testGetAllElementsInRDDWithVisibilityFilteringApplied(final Graph graph,
                                                                       final GetRDDOfAllElements getRDD)
            throws OperationException, IOException, InterruptedException, AccumuloSecurityException, AccumuloException {
        final Set<Element> expectedElements = new HashSet<>();

        // Test with user with public visibility
        getElementsWithVisibilities()
                .stream()
                .filter(e -> e.getProperty(TestPropertyNames.VISIBILITY).equals("public"))
                .forEach(expectedElements::add);
        RDD<Element> rdd = graph.execute(getRDD, USER_WITH_PUBLIC);
        if (rdd == null) {
            fail("No RDD returned");
        }
        final Set<Element> results = new HashSet<>();
        Element[] returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        // Test with user with public and private visibility
        getElementsWithVisibilities().forEach(expectedElements::add);
        rdd = graph.execute(getRDD, USER_WITH_PUBLIC_AND_PRIVATE);
        if (rdd == null) {
            fail("No RDD returned");
        }
        results.clear();
        returnedElements = (Element[]) rdd.collect();
        for (int i = 0; i < returnedElements.length; i++) {
            results.add(returnedElements[i]);
        }
        assertEquals(expectedElements, results);

        // Test with user with no visibilities
        rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }
        results.clear();
        returnedElements = (Element[]) rdd.collect();
        assertEquals(0, returnedElements.length);

        getRDD.getSparkSession().stop();
    }

    private void testGetAllElementsInRDDWithValidationApplied(final Graph graph, final GetRDDOfAllElements getRDD)
            throws InterruptedException, IOException, OperationException {
        // Sleep for 1 second to give chance for Entity A to age off
        Thread.sleep(1000L);

        final RDD<Element> rdd = graph.execute(getRDD, USER);
        if (rdd == null) {
            fail("No RDD returned");
        }

        // Should get Entity B but not Entity A
        final Element[] returnedElements = (Element[]) rdd.collect();
        assertEquals(1, returnedElements.length);

        assertEquals(entityRetainedAfterValidation, returnedElements[0]);

        getRDD.getSparkSession().stop();
    }

    private Graph getGraphForMockAccumulo() throws OperationException {
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(GRAPH_ID)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elements.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getClass().getResourceAsStream("/store.properties"))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElements()).build(), user);
        return graph;
    }

    private Graph getGraphForMockAccumuloWithVisibility() throws OperationException {
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(GRAPH_ID)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsWithVisibility.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getClass().getResourceAsStream("/store.properties"))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsWithVisibilities()).build(), user);
        return graph;
    }

    private Graph getGraphForMockAccumuloForValidationChecking() throws OperationException {
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(GRAPH_ID)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/typesForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getClass().getResourceAsStream("/store.properties"))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsForValidationChecking()).build(), user);
        return graph;
    }

    private Graph getGraphForDirectRDD(final String tableName) throws InterruptedException, AccumuloException, AccumuloSecurityException,
            IOException, OperationException, TableNotFoundException {
        final MiniAccumuloCluster cluster = MiniAccumuloClusterProvider.getMiniAccumuloCluster();
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(tableName)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elements.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getAccumuloProperties(cluster))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElements()).build(), user);
        cluster.getConnector("root", MiniAccumuloClusterProvider.PASSWORD)
                .tableOperations()
                .compact(tableName, new CompactionConfig());
        Thread.sleep(1000L);
        return graph;
    }

    private Graph getGraphForDirectRDDWithVisibility(final String tableName) throws InterruptedException, AccumuloException,
            AccumuloSecurityException, IOException, OperationException, TableNotFoundException {
        final MiniAccumuloCluster cluster = MiniAccumuloClusterProvider.getMiniAccumuloCluster();
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(tableName)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsWithVisibility.json"))
                .addSchema(getClass().getResourceAsStream("/schema/types.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getAccumuloProperties(cluster))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsWithVisibilities()).build(), user);
        cluster.getConnector("root", MiniAccumuloClusterProvider.PASSWORD)
                .tableOperations()
                .compact(tableName, new CompactionConfig());
        Thread.sleep(1000L);
        return graph;
    }

    private Graph getGraphForDirectRDDForValidationChecking(final String tableName) throws InterruptedException,
            AccumuloException, AccumuloSecurityException, IOException, OperationException, TableNotFoundException {
        final MiniAccumuloCluster cluster = MiniAccumuloClusterProvider.getMiniAccumuloCluster();
        final Graph graph = new Graph.Builder()
                .config(new GraphConfig.Builder()
                        .graphId(tableName)
                        .build())
                .addSchema(getClass().getResourceAsStream("/schema/elementsForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/typesForValidationChecking.json"))
                .addSchema(getClass().getResourceAsStream("/schema/serialisation.json"))
                .storeProperties(getAccumuloProperties(cluster))
                .build();
        final User user = new User();
        graph.execute(new AddElements.Builder().input(getElementsForValidationChecking()).build(), user);
        cluster.getConnector("root", MiniAccumuloClusterProvider.PASSWORD)
                .tableOperations()
                .compact(tableName, new CompactionConfig());
        Thread.sleep(1000L);
        return graph;
    }

    private List<Element> getElements() {
        final List<Element> elements = new ArrayList<>();
        for (int i = 0; i < 1; i++) {
            final Entity entity = new Entity.Builder()
                    .group(TestGroups.ENTITY)
                    .vertex("" + i)
                    .build();

            final Edge edge1 = new Edge.Builder()
                    .group(TestGroups.EDGE)
                    .source("" + i)
                    .dest("B")
                    .directed(false)
                    .property(TestPropertyNames.COUNT, 2)
                    .build();

            final Edge edge2 = new Edge.Builder()
                    .group(TestGroups.EDGE)
                    .source("" + i)
                    .dest("C")
                    .directed(false)
                    .property(TestPropertyNames.COUNT, 4)
                    .build();

            elements.add(edge1);
            elements.add(edge2);
            elements.add(entity);
        }
        return elements;
    }

    private List<Element> getElementsWithVisibilities() {
        return getElements().stream()
                .map(e -> {
                    if (e.getGroup().equals(TestGroups.ENTITY)) {
                        e.putProperty(TestPropertyNames.VISIBILITY, "public");
                    } else if (((Edge) e).getDestination().equals("B")) {
                        e.putProperty(TestPropertyNames.VISIBILITY, "private");
                    } else {
                        e.putProperty(TestPropertyNames.VISIBILITY, "public");
                    }
                    return e;
                })
                .collect(Collectors.toList());
    }

    private List<Element> getElementsForValidationChecking() {
        final List<Element> elements = new ArrayList<>();
        final Entity entity1 = new Entity.Builder()
                .group(TestGroups.ENTITY)
                .vertex("A")
                .property("timestamp", System.currentTimeMillis())
                .build();
        final Entity entity2 = new Entity.Builder()
                .group(TestGroups.ENTITY)
                .vertex("B")
                .property("timestamp", System.currentTimeMillis() + 1000000L)
                .build();
        entityRetainedAfterValidation = entity2;
        elements.add(entity1);
        elements.add(entity2);
        return elements;
    }

    private GetRDDOfAllElements getOperation(final String appName) throws IOException {
        // Create Hadoop configuration and serialise to a string
        final Configuration configuration = new Configuration();
        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        configuration.write(new DataOutputStream(baos));
        final String configurationString = new String(baos.toByteArray(), CommonConstants.UTF_8);

        // Check get correct elements
        final GetRDDOfAllElements rddQuery = new GetRDDOfAllElements.Builder()
                .sparkSession(getSparkSession(appName))
                .build();
        rddQuery.addOption(AbstractGetRDDHandler.HADOOP_CONFIGURATION_KEY, configurationString);
        return rddQuery;
    }

    private GetRDDOfAllElements getOperationWithDirectRDDOption(final String appName) throws IOException {
        final GetRDDOfAllElements op = getOperation(appName);
        op.addOption(AbstractGetRDDHandler.USE_RFILE_READER_RDD, "true");
        return op;
    }

    private AccumuloProperties getAccumuloProperties(final MiniAccumuloCluster cluster) {
        final AccumuloProperties properties = new AccumuloProperties();
        properties.setStoreClass(AccumuloStore.class);
        properties.setInstance(cluster.getInstanceName());
        properties.setZookeepers(cluster.getZooKeepers());
        properties.setUser(MiniAccumuloClusterProvider.USER);
        properties.setPassword(MiniAccumuloClusterProvider.PASSWORD);
        properties.setOperationDeclarationPaths("sparkAccumuloOperationsDeclarations.json");
        return properties;
    }

    private SparkSession getSparkSession(final String appName) {
        final SparkSession sparkSession = SparkSession.builder()
                .master("local")
                .appName(appName)
                .config(SparkConstants.SERIALIZER, SparkConstants.DEFAULT_SERIALIZER)
                .config(SparkConstants.KRYO_REGISTRATOR, SparkConstants.DEFAULT_KRYO_REGISTRATOR)
                .config(SparkConstants.DRIVER_ALLOW_MULTIPLE_CONTEXTS, true)
                .getOrCreate();
        return sparkSession;
    }
}
{'fraction_non_alphanumeric': 0.09914742118330444, 'fraction_numerical': 0.0029512343654846114, 'mean_word_length': 2.692380870016432, 'pattern_counts': {'":': 0, '<': 60, '<?xml version=': 0, '>': 60, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3963363', 'n_tokens_mistral': 11213, 'n_tokens_neox': 10733, 'n_words': 2027}