package com.tinkerpop.blueprints.util.wrappers.batch;
import com.tinkerpop.blueprints.BaseTest;
import com.tinkerpop.blueprints.Direction;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.Features;
import com.tinkerpop.blueprints.Graph;
import com.tinkerpop.blueprints.GraphQuery;
import com.tinkerpop.blueprints.TransactionalGraph;
import com.tinkerpop.blueprints.Vertex;
import com.tinkerpop.blueprints.impls.tg.IgnoreIdTinkerGraph;
import com.tinkerpop.blueprints.impls.tg.MockTransactionalGraph;
import com.tinkerpop.blueprints.impls.tg.TinkerGraph;
import junit.framework.TestCase;
import java.util.Random;
/**
* Tests {@link BatchGraph} by creating a variable length chain and verifying that the chain is correctly inserted into the wrapped TinkerGraph.
* <br />
* Tests the various different Vertex caches and different length of chains.
* <br />
*
* @author Matthias Broecheler (http://www.matthiasb.com)
*/
public class BatchGraphTest extends TestCase {
private static final String UID = "uid";
private static final String vertexIDKey = "vid";
private static final String edgeIDKey = "eid";
private static boolean assignKeys = false;
private static boolean ignoreIDs = false;
public void testNumberIdLoading() {
loadingTest(5000, 100, VertexIDType.NUMBER, new NumberLoadingFactory());
loadingTest(200000, 10000, VertexIDType.NUMBER, new NumberLoadingFactory());
assignKeys = true;
loadingTest(5000, 100, VertexIDType.NUMBER, new NumberLoadingFactory());
loadingTest(50000, 10000, VertexIDType.NUMBER, new NumberLoadingFactory());
assignKeys = false;
ignoreIDs = true;
loadingTest(5000, 100, VertexIDType.NUMBER, new NumberLoadingFactory());
loadingTest(50000, 10000, VertexIDType.NUMBER, new NumberLoadingFactory());
ignoreIDs = false;
}
public void testObjectIdLoading() {
loadingTest(5000, 100, VertexIDType.OBJECT, new StringLoadingFactory());
loadingTest(200000, 10000, VertexIDType.OBJECT, new StringLoadingFactory());
}
public void testStringIdLoading() {
loadingTest(5000, 100, VertexIDType.STRING, new StringLoadingFactory());
loadingTest(200000, 10000, VertexIDType.STRING, new StringLoadingFactory());
}
public void testURLIdLoading() {
loadingTest(5000, 100, VertexIDType.URL, new URLLoadingFactory());
loadingTest(200000, 10000, VertexIDType.URL, new URLLoadingFactory());
}
public void testQuadLoading() {
int numEdges = 10000;
String[][] quads = generateQuads(100, numEdges, new String[]{"knows", "friend"});
TinkerGraph graph = new TinkerGraph();
BatchGraph bgraph = new BatchGraph(new WritethroughGraph(graph), VertexIDType.STRING, 1000);
for (String[] quad : quads) {
Vertex[] vertices = new Vertex[2];
for (int i = 0; i < 2; i++) {
vertices[i] = bgraph.getVertex(quad[i]);
if (vertices[i] == null) vertices[i] = bgraph.addVertex(quad[i]);
}
Edge edge = bgraph.addEdge(null, vertices[0], vertices[1], quad[2]);
edge.setProperty("annotation", quad[3]);
}
assertEquals(numEdges, BaseTest.count(graph.getEdges()));
bgraph.shutdown();
}
public void testLoadingWithExisting1() {
int numEdges = 1000;
String[][] quads = generateQuads(100, numEdges, new String[]{"knows", "friend"});
TinkerGraph tg = new TinkerGraph();
BatchGraph bg = new BatchGraph(new WritethroughGraph(tg), VertexIDType.STRING, 100);
bg.setLoadingFromScratch(false);
Graph graph = null;
int counter = 0;
for (String[] quad : quads) {
if (counter < numEdges / 2) graph = tg;
else graph = bg;
Vertex[] vertices = new Vertex[2];
for (int i = 0; i < 2; i++) {
vertices[i] = graph.getVertex(quad[i]);
if (vertices[i] == null) vertices[i] = graph.addVertex(quad[i]);
}
Edge edge = graph.addEdge(null, vertices[0], vertices[1], quad[2]);
edge.setProperty("annotation", quad[3]);
counter++;
}
assertEquals(numEdges, BaseTest.count(tg.getEdges()));
bg.shutdown();
}
public void testLoadingWithExisting2() {
int numEdges = 1000;
String[][] quads = generateQuads(100, numEdges, new String[]{"knows", "friend"});
TinkerGraph tg = new IgnoreIdTinkerGraph();
BatchGraph bg = new BatchGraph(new WritethroughGraph(tg), VertexIDType.STRING, 100);
try {
bg.setLoadingFromScratch(false);
fail();
} catch (IllegalStateException e) {
}
bg.setVertexIdKey("uid");
bg.setLoadingFromScratch(false);
try {
bg.setVertexIdKey(null);
fail();
} catch (IllegalStateException e) {
}
Graph graph = null;
int counter = 0;
for (String[] quad : quads) {
if (counter < numEdges / 2) graph = tg;
else graph = bg;
Vertex[] vertices = new Vertex[2];
for (int i = 0; i < 2; i++) {
vertices[i] = graph.getVertex(quad[i]);
if (vertices[i] == null) vertices[i] = graph.addVertex(quad[i]);
}
Edge edge = graph.addEdge(null, vertices[0], vertices[1], quad[2]);
edge.setProperty("annotation", quad[3]);
counter++;
}
assertEquals(numEdges, BaseTest.count(tg.getEdges()));
bg.shutdown();
}
public static String[][] generateQuads(int numVertices, int numEdges, String[] labels) {
Random random = new Random();
String[][] edges = new String[numEdges][4];
for (int i = 0; i < numEdges; i++) {
edges[i][0] = "v" + random.nextInt(numVertices) + 1;
edges[i][1] = "v" + random.nextInt(numVertices) + 1;
edges[i][2] = labels[random.nextInt(labels.length)];
edges[i][3] = "" + random.nextInt();
}
return edges;
}
public void loadingTest(int total, int bufferSize, VertexIDType type, LoadingFactory ids) {
final VertexEdgeCounter counter = new VertexEdgeCounter();
MockTransactionalGraph tgraph = null;
if (ignoreIDs) {
tgraph = new MockTransactionalGraph(new IgnoreIdTinkerGraph());
} else {
tgraph = new MockTransactionalGraph(new TinkerGraph());
}
BLGraph graph = new BLGraph(tgraph, counter, ids);
BatchGraph<BLGraph> loader = new BatchGraph<BLGraph>(graph, type, bufferSize);
if (assignKeys) {
loader.setVertexIdKey(vertexIDKey);
loader.setEdgeIdKey(edgeIDKey);
}
//Create a chain
int chainLength = total;
Vertex previous = null;
for (int i = 0; i <= chainLength; i++) {
Vertex next = loader.addVertex(ids.getVertexID(i));
next.setProperty(UID, i);
counter.numVertices++;
counter.totalVertices++;
if (previous != null) {
Edge e = loader.addEdge(ids.getEdgeID(i), loader.getVertex(previous.getId()), loader.getVertex(next.getId()), "next");
e.setProperty(UID, i);
counter.numEdges++;
}
previous = next;
}
loader.stopTransaction(TransactionalGraph.Conclusion.SUCCESS);
assertTrue(tgraph.allSuccessful());
loader.shutdown();
}
static class VertexEdgeCounter {
int numVertices = 0;
int numEdges = 0;
int totalVertices = 0;
}
static class BLGraph implements TransactionalGraph {
private static final int keepLast = 10;
private final VertexEdgeCounter counter;
private boolean first = true;
private final LoadingFactory ids;
private final TransactionalGraph graph;
BLGraph(TransactionalGraph graph, final VertexEdgeCounter counter, LoadingFactory ids) {
this.graph = graph;
this.counter = counter;
this.ids = ids;
}
private static final Object parseID(Object id) {
if (id instanceof String) {
try {
return Integer.parseInt((String) id);
} catch (NumberFormatException e) {
return id;
}
} else return id;
}
@Override
public void commit() {
graph.commit();
verifyCounts();
}
@Override
public void rollback() {
graph.rollback();
verifyCounts();
}
@Override
public void stopTransaction(Conclusion conclusion) {
if (Conclusion.SUCCESS == conclusion)
commit();
else
rollback();
}
private void verifyCounts() {
//System.out.println("Committed (vertices/edges): " + counter.numVertices + " / " + counter.numEdges);
assertEquals(counter.numVertices, BaseTest.count(graph.getVertices()) - (first ? 0 : keepLast));
assertEquals(counter.numEdges, BaseTest.count(graph.getEdges()));
for (Edge e : getEdges()) {
int id = ((Number) e.getProperty(UID)).intValue();
if (!ignoreIDs) {
assertEquals(ids.getEdgeID(id), parseID(e.getId()));
}
assertEquals(1, (Integer) e.getVertex(Direction.IN).getProperty(UID) - (Integer) e.getVertex(Direction.OUT).getProperty(UID));
if (assignKeys) {
assertEquals(ids.getEdgeID(id), e.getProperty(edgeIDKey));
}
}
for (Vertex v : getVertices()) {
int id = ((Number) v.getProperty(UID)).intValue();
if (!ignoreIDs) {
assertEquals(ids.getVertexID(id), parseID(v.getId()));
}
assertTrue(2 >= BaseTest.count(v.getEdges(Direction.BOTH)));
assertTrue(1 >= BaseTest.count(v.getEdges(Direction.IN)));
assertTrue(1 >= BaseTest.count(v.getEdges(Direction.OUT)));
if (assignKeys) {
assertEquals(ids.getVertexID(id), v.getProperty(vertexIDKey));
}
}
for (Vertex v : getVertices()) {
int id = ((Number) v.getProperty(UID)).intValue();
if (id < counter.totalVertices - keepLast) {
removeVertex(v);
}
}
for (Edge e : getEdges()) removeEdge(e);
assertEquals(keepLast, BaseTest.count(graph.getVertices()));
counter.numVertices = 0;
counter.numEdges = 0;
first = false;
//System.out.println("------");
}
@Override
public Features getFeatures() {
return graph.getFeatures();
}
@Override
public Vertex addVertex(Object id) {
return graph.addVertex(id);
}
@Override
public Vertex getVertex(Object id) {
return graph.getVertex(id);
}
@Override
public void removeVertex(Vertex vertex) {
graph.removeVertex(vertex);
}
@Override
public Iterable<Vertex> getVertices() {
return graph.getVertices();
}
@Override
public Iterable<Vertex> getVertices(String key, Object value) {
return graph.getVertices(key, value);
}
@Override
public Edge addEdge(Object id, Vertex outVertex, Vertex inVertex, String label) {
return graph.addEdge(id, outVertex, inVertex, label);
}
@Override
public Edge getEdge(Object id) {
return graph.getEdge(id);
}
@Override
public void removeEdge(Edge edge) {
graph.removeEdge(edge);
}
@Override
public Iterable<Edge> getEdges() {
return graph.getEdges();
}
@Override
public Iterable<Edge> getEdges(String key, Object value) {
return graph.getEdges(key, value);
}
@Override
public void shutdown() {
graph.shutdown();
}
@Override
public GraphQuery query() {
return graph.query();
}
}
interface LoadingFactory {
public Object getVertexID(int id);
public Object getEdgeID(int id);
}
class StringLoadingFactory implements LoadingFactory {
@Override
public Object getVertexID(int id) {
return "V" + id;
}
@Override
public Object getEdgeID(int id) {
return "E" + id;
}
}
class NumberLoadingFactory implements LoadingFactory {
@Override
public Object getVertexID(int id) {
return Integer.valueOf(id * 2);
}
@Override
public Object getEdgeID(int id) {
return Integer.valueOf(id * 2 + 1);
}
}
class URLLoadingFactory implements LoadingFactory {
@Override
public Object getVertexID(int id) {
return "http://www.tinkerpop.com/rdf/ns/vertex/" + id;
}
@Override
public Object getEdgeID(int id) {
return "http://www.tinkerpop.com/rdf/ns/edge#" + id;
}
}
}