Package com.tinkerpop.gremlin.groovy.engine

Source Code of com.tinkerpop.gremlin.groovy.engine.GremlinExecutorTest

package com.tinkerpop.gremlin.groovy.engine;

import com.tinkerpop.gremlin.groovy.jsr223.GremlinGroovyScriptEngineTest;
import com.tinkerpop.gremlin.structure.Graph;
import com.tinkerpop.gremlin.structure.io.graphson.GraphSONResourceAccess;
import com.tinkerpop.gremlin.structure.io.kryo.KryoResourceAccess;
import org.junit.Ignore;
import org.junit.Test;
import org.kohsuke.groovy.sandbox.GroovyInterceptor;

import javax.script.Bindings;
import javax.script.SimpleBindings;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

/**
* @author Stephen Mallette (http://stephen.genoprime.com)
*/
public class GremlinExecutorTest {
    public static Map<String, String> PATHS = new HashMap<>();

    static {
        try {
            final List<String> groovyScriptResources = Arrays.asList("GremlinExecutorInit.groovy");
            for (final String fileName : groovyScriptResources) {
                PATHS.put(fileName, generateTempFile(GremlinExecutorTest.class, fileName));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Test
    public void shouldEvalScript() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().create();
        assertEquals(2, gremlinExecutor.eval("1+1").get());
    }

    @Test
    public void shouldEvalMultipleScripts() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().create();
        assertEquals(2, gremlinExecutor.eval("1+1").get());
        assertEquals(3, gremlinExecutor.eval("1+2").get());
        assertEquals(4, gremlinExecutor.eval("1+3").get());
        assertEquals(5, gremlinExecutor.eval("1+4").get());
        assertEquals(6, gremlinExecutor.eval("1+5").get());
        assertEquals(7, gremlinExecutor.eval("1+6").get());
    }

    @Test
    public void shouldEvalScriptWithBindings() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().create();
        final Bindings b = new SimpleBindings();
        b.put("x", 1);
        assertEquals(2, gremlinExecutor.eval("1+x", b).get());
    }

    @Test
    public void shouldEvalScriptWithGlobalBindings() throws Exception {
        final Bindings b = new SimpleBindings();
        b.put("x", 1);
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().globalBindings(b).create();
        assertEquals(2, gremlinExecutor.eval("1+x").get());
    }

    @Test
    public void shouldEvalScriptWithGlobalAndLocalBindings() throws Exception {
        final Bindings g = new SimpleBindings();
        g.put("x", 1);
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().globalBindings(g).create();
        final Bindings b = new SimpleBindings();
        b.put("y", 1);
        assertEquals(2, gremlinExecutor.eval("y+x", b).get());
    }

    @Test
    public void shouldEvalScriptWithLocalOverridingGlobalBindings() throws Exception {
        final Bindings g = new SimpleBindings();
        g.put("x", 1);
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().globalBindings(g).create();
        final Bindings b = new SimpleBindings();
        b.put("x", 10);
        assertEquals(11, gremlinExecutor.eval("x+1", b).get());
    }

    @Test
    public void shouldTimeoutScript() throws Exception {
        final AtomicBoolean timeoutCalled = new AtomicBoolean(false);
        final AtomicBoolean successCalled = new AtomicBoolean(false);
        final AtomicBoolean failureCalled = new AtomicBoolean(false);
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .scriptEvaluationTimeout(500)
                .afterFailure((b, e) -> failureCalled.set(true))
                .afterSuccess((b) -> successCalled.set(true))
                .afterTimeout((b) -> timeoutCalled.set(true)).create();
        try {
            gremlinExecutor.eval("Thread.sleep(1000);10").get();
            fail();
        } catch (Exception ex) {

        }

        // need to wait long enough for the script to complete
        Thread.sleep(750);

        assertTrue(timeoutCalled.get());
        assertFalse(successCalled.get());
        assertFalse(failureCalled.get());
    }

    @Test
    public void shouldCallFail() throws Exception {
        final AtomicBoolean timeoutCalled = new AtomicBoolean(false);
        final AtomicBoolean successCalled = new AtomicBoolean(false);
        final AtomicBoolean failureCalled = new AtomicBoolean(false);
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .afterFailure((b, e) -> failureCalled.set(true))
                .afterSuccess((b) -> successCalled.set(true))
                .afterTimeout((b) -> timeoutCalled.set(true)).create();
        try {
            gremlinExecutor.eval("10/0").get();
            fail();
        } catch (Exception ex) {

        }

        // need to wait long enough for the script to complete
        Thread.sleep(750);

        assertFalse(timeoutCalled.get());
        assertFalse(successCalled.get());
        assertTrue(failureCalled.get());
    }

    @Test
    public void shouldCallSuccess() throws Exception {
        final AtomicBoolean timeoutCalled = new AtomicBoolean(false);
        final AtomicBoolean successCalled = new AtomicBoolean(false);
        final AtomicBoolean failureCalled = new AtomicBoolean(false);
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .afterFailure((b, e) -> failureCalled.set(true))
                .afterSuccess((b) -> successCalled.set(true))
                .afterTimeout((b) -> timeoutCalled.set(true)).create();
        assertEquals(2, gremlinExecutor.eval("1+1").get());

        // need to wait long enough for the script to complete
        Thread.sleep(750);

        assertFalse(timeoutCalled.get());
        assertTrue(successCalled.get());
        assertFalse(failureCalled.get());
    }

    @Test
    public void shouldEvalInMultipleThreads() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().create();

        final CyclicBarrier barrier = new CyclicBarrier(2);
        final AtomicInteger i1 = new AtomicInteger(0);
        final AtomicBoolean b1 = new AtomicBoolean(false);
        final Thread t1 = new Thread(() -> {
            try {
                barrier.await();
                i1.set((Integer) gremlinExecutor.eval("1+1").get());
            } catch (Exception ex) {
                b1.set(true);
            }
        });

        final AtomicInteger i2 = new AtomicInteger(0);
        final AtomicBoolean b2 = new AtomicBoolean(false);
        final Thread t2 = new Thread(() -> {
            try {
                barrier.await();
                i2.set((Integer) gremlinExecutor.eval("1+1").get());
            } catch (Exception ex) {
                b2.set(true);
            }
        });

        t1.start();
        t2.start();

        t1.join();
        t2.join();

        assertEquals(2, i1.get());
        assertEquals(2, i2.get());
        assertFalse(b1.get());
        assertFalse(b2.get());
    }

    @Test
    public void shouldNotExhaustThreads() throws Exception {
        // this is not representative of how the GremlinExecutor should be configured.  A single thread executor
        // shared will create odd behaviors, but it's good for this test.
        final ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .executorService(executorService)
                .scheduledExecutorService(executorService).create();

        final AtomicInteger count = new AtomicInteger(0);
        assertTrue(IntStream.range(0, 1000).mapToObj(i -> gremlinExecutor.eval("1+1")).allMatch(f -> {
            try {
                return (Integer) f.get() == 2;
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            } finally {
                count.incrementAndGet();
            }
        }));

        assertEquals(1000, count.intValue());
    }

    @Test
    public void shouldFailUntilImportExecutes() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build().create();

        final Set<String> imports = new HashSet<String>() {{
            add("import java.awt.Color");
        }};

        final AtomicInteger successes = new AtomicInteger(0);
        final AtomicInteger failures = new AtomicInteger(0);

        // issue 1000 scripts in one thread using a class that isn't imported.  this will result in failure.
        // while that thread is running start a new thread that issues an addImports to include that class.
        // this should block further evals in the first thread until the import is complete at which point
        // evals in the first thread will resume and start to succeed
        final Thread t1 = new Thread(() ->
                IntStream.range(0, 1000).mapToObj(i -> gremlinExecutor.eval("Color.BLACK"))
                        .forEach(f -> {
                            f.exceptionally(t -> failures.incrementAndGet()).join();
                            if (!f.isCompletedExceptionally())
                                successes.incrementAndGet();
                        })
        );

        final Thread t2 = new Thread(() -> {
            while (failures.get() < 500) {
            }
            gremlinExecutor.getScriptEngines().addImports(imports);
        });

        t1.start();
        t2.start();

        t1.join();
        t2.join();

        assertTrue(successes.intValue() > 0);
        assertTrue(failures.intValue() >= 500);
    }

    @Test
    public void shouldInitializeWithScript() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .addEngineSettings("gremlin-groovy",
                        Collections.emptyList(),
                        Collections.emptyList(),
                        Arrays.asList(PATHS.get("GremlinExecutorInit.groovy")),
                        Collections.emptyMap())
                .create();

        assertEquals(2, gremlinExecutor.eval("sum(1,1)").get());
    }

    @Test
    public void shouldSecureAll() throws Exception {
        GroovyInterceptor.getApplicableInterceptors().forEach(GroovyInterceptor::unregister);
        final Map<String,Object> config = new HashMap<>();
        config.put("sandbox", GremlinGroovyScriptEngineTest.DenyAll.class.getName());
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .addEngineSettings("gremlin-groovy",
                        Collections.emptyList(),
                        Collections.emptyList(),
                        Arrays.asList(PATHS.get("GremlinExecutorInit.groovy")),
                        config)
                .create();
        try {
            gremlinExecutor.eval("c = new java.awt.Color(255, 255, 255)").get();
            fail("Should have failed security");
        } catch (Exception se) {
            assertEquals(SecurityException.class, se.getCause().getCause().getCause().getCause().getClass());
        } finally {
            gremlinExecutor.close();
        }
    }

    @Test
    public void shouldSecureSome() throws Exception {
        GroovyInterceptor.getApplicableInterceptors().forEach(GroovyInterceptor::unregister);
        final Map<String,Object> config = new HashMap<>();
        config.put("sandbox", GremlinGroovyScriptEngineTest.AllowSome.class.getName());
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .addEngineSettings("gremlin-groovy",
                        Collections.emptyList(),
                        Collections.emptyList(),
                        Arrays.asList(PATHS.get("GremlinExecutorInit.groovy")),
                        config)
                .create();
        try {
            gremlinExecutor.eval("c = 'new java.awt.Color(255, 255, 255)'").get();
            fail("Should have failed security");
        } catch (Exception se) {
            assertEquals(SecurityException.class, se.getCause().getCause().getCause().getCause().getClass());
        }

        try {
            final java.awt.Color c = (java.awt.Color) gremlinExecutor.eval("g = new java.awt.Color(255, 255, 255)").get();
            assertEquals(java.awt.Color.class, c.getClass());
        } catch (Exception ignored) {
            fail("Should not have tossed an exception");
        } finally {
            gremlinExecutor.close();
        }
    }

    @Test
    public void shouldInitializeWithScriptAndWorkAfterReset() throws Exception {
        final GremlinExecutor gremlinExecutor = GremlinExecutor.build()
                .addEngineSettings("gremlin-groovy",
                        Collections.emptyList(),
                        Collections.emptyList(),
                        Arrays.asList(PATHS.get("GremlinExecutorInit.groovy")),
                        Collections.emptyMap())
                .create();

        assertEquals(2, gremlinExecutor.eval("sum(1,1)").get());

        gremlinExecutor.getScriptEngines().reset();

        assertEquals(2, gremlinExecutor.eval("sum(1,1)").get());
    }

    public static String generateTempFile(final Class resourceClass, final String fileName) throws IOException {
        final File temp = File.createTempFile(fileName, ".groovy");
        final FileOutputStream outputStream = new FileOutputStream(temp);
        int data;
        final InputStream inputStream = resourceClass.getResourceAsStream(fileName);
        while ((data = inputStream.read()) != -1) {
            outputStream.write(data);
        }
        outputStream.close();
        inputStream.close();
        return temp.getPath();
    }
}
TOP

Related Classes of com.tinkerpop.gremlin.groovy.engine.GremlinExecutorTest

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.