package ai.grakn.test.benchmark;

import ai.grakn.Grakn;
import ai.grakn.GraknSession;
import ai.grakn.GraknTx;
import ai.grakn.GraknTxType;
import ai.grakn.Keyspace;
import ai.grakn.concept.EntityType;
import ai.grakn.graql.GetQuery;
import ai.grakn.graql.Match;
import ai.grakn.graql.admin.Answer;
import ai.grakn.test.EngineContext;
import ai.grakn.util.SampleKBLoader;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.TearDown;

import java.util.Optional;

import static ai.grakn.graql.Graql.var;


public class MatchBenchmark extends BenchmarkTest {

    private static final Keyspace KEYSPACE = SampleKBLoader.randomKeyspace();
    private static final String BENCHMARK_ENTITYTYPE = "benchmarkEntitytype";

    private EngineContext engine;
    private GraknSession session;
    private GraknTx graph;

    @Setup
    public void setup() throws Throwable {
        engine = EngineContext.createWithInMemoryRedis();
        engine.before();
        session = Grakn.session(engine.uri(), KEYSPACE);
        GraknTx graphEntity = session.open(GraknTxType.WRITE);
        EntityType entityType = graphEntity.putEntityType("benchmarkEntitytype");
        entityType.addEntity();
        graphEntity.commit();
        graph = session.open(GraknTxType.WRITE);
    }

    @TearDown
    public void tearDown() {
        graph.close();
        session.close();
        engine.after();
    }

    @Benchmark
    public void match() {
        Match match = graph.graql().match(var("x").isa(BENCHMARK_ENTITYTYPE));
        GetQuery answers = match.get();
        Optional<Answer> first = answers.stream().findFirst();
        first.get();
    }
}

@@ -160,16 +160,55 @@
 yspace;%0A
+import ai.grakn.concept.AttributeType;%0A
 import a
@@ -231,24 +231,24 @@
 EntityType;%0A
-
 import ai.gr
@@ -784,24 +784,25 @@
 HMARK_ENTITY
+_
 TYPE = %22benc
@@ -808,25 +808,110 @@
 chmarkEntity
-t
+Type%22;%0A    private static final String BENCHMARK_ATTRIBUTE_TYPE = %22benchmarkAttributeT
 ype%22;%0A%0A    p
@@ -1320,63 +1320,431 @@
 ype(
-%22benchmarkEntitytype%22);%0A        entityType.addEntity();
+BENCHMARK_ENTITY_TYPE);%0A        AttributeType%3CString%3E attributeType =%0A                graphEntity.putAttributeType(BENCHMARK_ATTRIBUTE_TYPE, AttributeType.DataType.STRING);%0A        entityType.attribute(attributeType);%0A%0A        for (int i = 0; i %3C 100; i++) %7B%0A            for (int j = 0; j %3C 100; j++) %7B%0A                entityType.addEntity().attribute(attributeType.putAttribute(String.valueOf(i)));%0A            %7D%0A        %7D
 %0A   
@@ -2035,16 +2035,58 @@
 tch(
-var(%22x%22)
+%0A                var(%22x%22)%0A                        
 .isa
@@ -2102,21 +2102,91 @@
 K_ENTITY
+_
 TYPE)
+%0A                        .has(BENCHMARK_ATTRIBUTE_TYPE, %220%22)%0A        
 );%0A     

Make match benchmark test a bit more similar to real use case (#2249)