Examples with Entity org.openstreetmap.osmosis.core.domain.v0_6.Entity used on opensource projects

Example 41 with Entity

use of org.openstreetmap.osmosis.core.domain.v0_6.Entity in project beam by apache.

the class DatastoreV1Test method testDeleteEntities.

/**
 * Test that entities with valid keys are transformed to delete mutations.
 */
@Test
public void testDeleteEntities() throws Exception {
    Key key = makeKey("bird", "finch").build();
    Entity entity = Entity.newBuilder().setKey(key).build();
    DeleteEntityFn deleteEntityFn = new DeleteEntityFn();
    Mutation expectedMutation = makeDelete(entity.getKey()).build();
    assertEquals(expectedMutation, deleteEntityFn.apply(entity));
}

Example 42 with Entity

use of org.openstreetmap.osmosis.core.domain.v0_6.Entity in project beam by apache.

the class DatastoreV1Test method readFnTest.

/**
 * Helper function to run a test reading from a {@link ReadFn}.
 */
private void readFnTest(int numEntities) throws Exception {
    // An empty query to read entities.
    Query query = Query.newBuilder().setLimit(Int32Value.newBuilder().setValue(numEntities)).build();
    // Use mockResponseForQuery to generate results.
    when(mockDatastore.runQuery(any(RunQueryRequest.class))).thenAnswer(invocationOnMock -> {
        Query q = ((RunQueryRequest) invocationOnMock.getArguments()[0]).getQuery();
        return mockResponseForQuery(q);
    });
    ReadFn readFn = new ReadFn(V_1_OPTIONS, mockDatastoreFactory);
    DoFnTester<Query, Entity> doFnTester = DoFnTester.of(readFn);
    /**
     * Although Datastore client is marked transient in {@link ReadFn}, when injected through mock
     * factory using a when clause for unit testing purposes, it is not serializable because it
     * doesn't have a no-arg constructor. Thus disabling the cloning to prevent the test object from
     * being serialized.
     */
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    List<Entity> entities = doFnTester.processBundle(query);
    int expectedNumCallsToRunQuery = (int) Math.ceil((double) numEntities / QUERY_BATCH_LIMIT);
    verify(mockDatastore, times(expectedNumCallsToRunQuery)).runQuery(any(RunQueryRequest.class));
    // Validate the number of results.
    assertEquals(numEntities, entities.size());
}

Example 43 with Entity

use of org.openstreetmap.osmosis.core.domain.v0_6.Entity in project beam by apache.

the class DatastoreV1Test method testDatatoreWriterFnWithLargeEntities.

/**
 * Tests {@link DatastoreWriterFn} with large entities that need to be split into more batches.
 */
@Test
public void testDatatoreWriterFnWithLargeEntities() throws Exception {
    List<Mutation> mutations = new ArrayList<>();
    int entitySize = 0;
    for (int i = 0; i < 12; ++i) {
        Entity entity = Entity.newBuilder().setKey(makeKey("key" + i, i + 1)).putProperties("long", makeValue(new String(new char[900_000])).setExcludeFromIndexes(true).build()).build();
        // Take the size of any one entity.
        entitySize = entity.getSerializedSize();
        mutations.add(makeUpsert(entity).build());
    }
    DatastoreWriterFn datastoreWriter = new DatastoreWriterFn(StaticValueProvider.of(PROJECT_ID), null, mockDatastoreFactory, new FakeWriteBatcher());
    DoFnTester<Mutation, Void> doFnTester = DoFnTester.of(datastoreWriter);
    doFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
    doFnTester.processBundle(mutations);
    // This test is over-specific currently; it requires that we split the 12 entity writes into 3
    // requests, but we only need each CommitRequest to be less than 10MB in size.
    int entitiesPerRpc = DATASTORE_BATCH_UPDATE_BYTES_LIMIT / entitySize;
    int start = 0;
    while (start < mutations.size()) {
        int end = Math.min(mutations.size(), start + entitiesPerRpc);
        CommitRequest.Builder commitRequest = CommitRequest.newBuilder();
        commitRequest.setMode(CommitRequest.Mode.NON_TRANSACTIONAL);
        commitRequest.addAllMutations(mutations.subList(start, end));
        // Verify all the batch requests were made with the expected mutations.
        verify(mockDatastore).commit(commitRequest.build());
        start = end;
    }
}

Example 44 with Entity

use of org.openstreetmap.osmosis.core.domain.v0_6.Entity in project beam by apache.

the class DatastoreV1Test method makeLatestTimestampResponse.

/**
 * Builds a response of the given timestamp.
 */
private static RunQueryResponse makeLatestTimestampResponse(long timestamp) {
    RunQueryResponse.Builder timestampResponse = RunQueryResponse.newBuilder();
    Entity.Builder entity = Entity.newBuilder();
    entity.setKey(makeKey("dummyKind", "dummyId"));
    entity.putProperties("timestamp", makeValue(new Date(timestamp * 1000)).build());
    EntityResult.Builder entityResult = EntityResult.newBuilder();
    entityResult.setEntity(entity);
    QueryResultBatch.Builder batch = QueryResultBatch.newBuilder();
    batch.addEntityResults(entityResult);
    timestampResponse.setBatch(batch);
    return timestampResponse.build();
}

Example 45 with Entity

use of org.openstreetmap.osmosis.core.domain.v0_6.Entity in project beam by apache.

the class V1ReadIT method writeEntitiesToDatastore.

// Creates entities and write them to datastore
private static void writeEntitiesToDatastore(V1TestOptions options, String project, String ancestor, long numEntities) throws Exception {
    Datastore datastore = getDatastore(options, project);
    // Write test entities to datastore
    V1TestWriter writer = new V1TestWriter(datastore, new UpsertMutationBuilder());
    Key ancestorKey = makeAncestorKey(options.getNamespace(), options.getKind(), ancestor);
    for (long i = 0; i < numEntities; i++) {
        Entity entity = makeEntity(i, ancestorKey, options.getKind(), options.getNamespace(), 0);
        writer.write(entity);
    }
    writer.close();
}