package com.tinkerpop.blueprints.generator;

import com.tinkerpop.blueprints.AbstractBlueprintsTest;
import com.tinkerpop.blueprints.BlueprintsStandardSuite;
import com.tinkerpop.blueprints.Graph;
import com.tinkerpop.blueprints.Property;
import com.tinkerpop.blueprints.util.StreamFactory;
import org.apache.commons.configuration.Configuration;
import org.junit.Test;
import org.junit.experimental.runners.Enclosed;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

import java.util.Arrays;

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

/**
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
@RunWith(Enclosed.class)
public class DistributionGeneratorTest {

    @RunWith(Parameterized.class)
    public static class DifferentDistributionTest extends AbstractBlueprintsTest {

        @Parameterized.Parameters(name = "{index}: test({0},{1})")
        public static Iterable<Object[]> data() {
            return Arrays.asList(new Object[][]{
                    {new NormalDistribution(2), null},
                    {new NormalDistribution(2), new NormalDistribution(5)},
                    {new PowerLawDistribution(2.1), null},
                    {new PowerLawDistribution(2.9), null},
                    {new PowerLawDistribution(3.9), null},
                    {new PowerLawDistribution(2.3), new PowerLawDistribution(2.8)}
            });
        }

        @Parameterized.Parameter(value = 0)
        public Distribution inDistribution;

        @Parameterized.Parameter(value = 1)
        public Distribution outDistribution;

        private static final int numberOfVertices = 100;

        @Test
        public void shouldGenerateRandomGraph() throws Exception {
            final DistributionGenerator generator = new DistributionGenerator("knows");
            distributionGeneratorTest(g, generator);

            final Configuration configuration = graphProvider.newGraphConfiguration("g1");
            final Graph g1 = graphProvider.openTestGraph(configuration);
            prepareGraph(g1);
            final DistributionGenerator generator1 = new DistributionGenerator("knows");
            distributionGeneratorTest(g1, generator1);

            assertNotEquals(StreamFactory.stream(g.query().edges()).count(), StreamFactory.stream(g1.query().edges()).count());

            graphProvider.clear(g1, configuration);
        }

        @Test
        public void shouldGenerateSameGraph() throws Exception {
            final DistributionGenerator generator = new DistributionGenerator("knows", null, ()->123456789l);
            distributionGeneratorTest(g, generator);

            final Configuration configuration = graphProvider.newGraphConfiguration("g1");
            final Graph g1 = graphProvider.openTestGraph(configuration);
            prepareGraph(g1);
            final DistributionGenerator generator1 = new DistributionGenerator("knows", null, ()->123456789l);
            distributionGeneratorTest(g1, generator1);

            assertEquals(StreamFactory.stream(g.query().edges()).count(), StreamFactory.stream(g1.query().edges()).count());

            graphProvider.clear(g1, configuration);
        }

        @Override
        protected void prepareGraph(final Graph g) throws Exception {
            final int numNodes = numberOfVertices;
            for (int i = 0; i < numNodes; i++) g.addVertex(Property.Key.ID, i);
        }

        private void distributionGeneratorTest(final Graph graph, final DistributionGenerator generator) {
            generator.setOutDistribution(inDistribution);
            if (outDistribution != null) generator.setOutDistribution(outDistribution);
            final int numEdges = generator.generate(graph, numberOfVertices * 10);
            tryCommit(graph, g -> assertEquals(numEdges, SizableIterable.sizeOf(g.query().edges())));
        }
    }
}

@@ -602,16 +602,59 @@
 tEquals;
+%0Aimport static org.junit.Assert.assertTrue;
 %0A%0A/**%0A *
@@ -860,16 +860,17 @@
 ribution
+s
 Test ext
@@ -3951,11 +3951,613 @@
  %7D%0A    %7D
+%0A%0A    public static class AnnotatorTest extends AbstractBlueprintsTest %7B%0A        @Test%0A        public void shouldAnnotateEdges() %7B%0A            final DistributionGenerator generator = new DistributionGenerator(%22knows%22, e-%3Ee.setProperty(%22data%22, %22test%22));%0A            final Distribution dist = new NormalDistribution(2);%0A            generator.setOutDistribution(dist);%0A            generator.setOutDistribution(dist);%0A            generator.generate(g, 100);%0A            tryCommit(g, g -%3E assertTrue(StreamFactory.stream(g.query().edges()).allMatch(e -%3E e.getValue(%22data%22).equals(%22test%22))));%0A        %7D%0A    %7D
 %0A%7D%0A

Add test for annotators.