Skip to content

Commit

Permalink
add bulk import test
Browse files Browse the repository at this point in the history
  • Loading branch information
vroldanbet committed May 3, 2024
1 parent dd86d42 commit a8ae90a
Showing 1 changed file with 83 additions and 0 deletions.
83 changes: 83 additions & 0 deletions src/intTest/java/V1ClientTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,12 @@
import java.util.Iterator;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.CountDownLatch;

import com.authzed.api.v1.*;
import com.authzed.grpcutil.BearerToken;

import io.grpc.stub.StreamObserver;
import org.junit.Test;

import io.grpc.ManagedChannel;
Expand Down Expand Up @@ -174,6 +176,87 @@ private static String writeRelationship(PermissionsServiceGrpc.PermissionsServic
return relResponse.getWrittenAt().getToken();
}


class BulkImportObserver implements StreamObserver<BulkImportRelationshipsResponse> {
final CountDownLatch done = new CountDownLatch(1);
private long loaded;

@Override
public void onNext(BulkImportRelationshipsResponse resp) {
loaded += resp.getNumLoaded();
}

@Override
public void onError(Throwable throwable) {
// TODO need to capture error so that blocking callsite is able to access it
System.out.println("onError");
done.countDown();
}

@Override
public void onCompleted() {
System.out.println("onCompleted");
done.countDown();
}

public void await() throws InterruptedException {
done.await();
}

public long loaded() {
return loaded;
}
};

@Test
public void testBulkImport() {

ManagedChannel channel = ManagedChannelBuilder.forTarget(target).usePlaintext().build();
String token = generateToken();
ExperimentalServiceGrpc.ExperimentalServiceStub experimentalService = ExperimentalServiceGrpc.
newStub(channel)
.withCallCredentials(new BearerToken(token));

BulkImportObserver responseObserver = new BulkImportObserver();
writeTestSchema(token, channel);
io.grpc.stub.StreamObserver<com.authzed.api.v1.BulkImportRelationshipsRequest>
observer = experimentalService.bulkImportRelationships(responseObserver);

for (int i = 0; i < 10; i++) {
BulkImportRelationshipsRequest req = BulkImportRelationshipsRequest.newBuilder()
.addRelationships(relationship("test/article", "java_test_" + i, "author", "test/user", "george")).build();
observer.onNext(req);
}
observer.onCompleted();

try {
responseObserver.await();
} catch (InterruptedException e) {
throw new RuntimeException(e);
}

assertEquals(10, responseObserver.loaded());
}

private static Relationship relationship(String resourceType, String resourceID, String relation, String subjectType, String subjectID) {
return Relationship.newBuilder()
.setResource(
ObjectReference.newBuilder()
.setObjectType(resourceType)
.setObjectId(resourceID)
.build())
.setRelation(relation)
.setSubject(
SubjectReference.newBuilder()
.setObject(
ObjectReference.newBuilder()
.setObjectType(subjectType)
.setObjectId(subjectID)
.build())
.build())
.build();
}

private static SchemaServiceGrpc.SchemaServiceBlockingStub writeTestSchema(String token, ManagedChannel channel) {
SchemaServiceGrpc.SchemaServiceBlockingStub schemaService = SchemaServiceGrpc.newBlockingStub(channel)
.withCallCredentials(new BearerToken(token));
Expand Down

0 comments on commit a8ae90a

Please sign in to comment.