Examples with EntityQueryServiceClient org.hypertrace.entity.query.service.client.EntityQueryServiceClient used on opensource projects

Example 1 with EntityQueryServiceClient

use of org.hypertrace.entity.query.service.client.EntityQueryServiceClient in project gateway-service by hypertrace.

the class EdsEntityUpdaterTest method testBulkUpdateEntitiesRequestForMultiValueOperation.

@Test
void testBulkUpdateEntitiesRequestForMultiValueOperation() {
    EntityQueryServiceClient mockEqsClient = mock(EntityQueryServiceClient.class);
    EdsEntityUpdater entityUpdater = new EdsEntityUpdater(mockEqsClient);
    BulkUpdateEntitiesRequest request = BulkUpdateEntitiesRequest.newBuilder().setEntityType("test-entity-type").addAllEntityIds(List.of("entity-id-1", "entity-id-2")).setOperation(BulkUpdateEntitiesRequestOperation.newBuilder().setMultiValuedAttributeOperation(MultiValuedAttributeOperation.newBuilder().setType(MultiValuedAttributeOperation.OperationType.OPERATION_TYPE_ADD).setAttribute(ColumnIdentifier.newBuilder().setColumnName("labels").build()).addAllValues(List.of(LiteralConstant.newBuilder().setValue(Value.newBuilder().setValueType(ValueType.STRING).setString("test-value").build()).build())).build()).build()).build();
    UpdateExecutionContext context = mock(UpdateExecutionContext.class);
    entityUpdater.bulkUpdateEntities(request, context);
    org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest expectedRequest = org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest.newBuilder().setEntityType("test-entity-type").addAllEntityIds(List.of("entity-id-1", "entity-id-2")).setOperation(org.hypertrace.entity.query.service.v1.BulkEntityArrayAttributeUpdateRequest.Operation.OPERATION_ADD).setAttribute(org.hypertrace.entity.query.service.v1.ColumnIdentifier.newBuilder().setColumnName("labels").build()).addAllValues(List.of(org.hypertrace.entity.query.service.v1.LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING).setString("test-value").build()).build())).build();
    verify(mockEqsClient).bulkUpdateEntityArrayAttribute(eq(expectedRequest), any());
}

Example 2 with EntityQueryServiceClient

use of org.hypertrace.entity.query.service.client.EntityQueryServiceClient in project gateway-service by hypertrace.

the class EdsEntityUpdaterTest method nonExistentEntity.

@Test
public void nonExistentEntity() {
    EntityQueryServiceClient mockEqsClient = mock(EntityQueryServiceClient.class);
    EdsEntityUpdater entityUpdater = new EdsEntityUpdater(mockEqsClient);
    // mock empty ResultSetChunk
    UpdateEntityRequest updateEntityRequest = createResolveTestRequest("non-existent-id");
    ResultSetChunk emptyResult = ResultSetChunk.newBuilder().build();
    when(mockEqsClient.update(any(), any())).thenReturn(List.of(emptyResult).iterator());
    UpdateExecutionContext updateExecutionContext = new UpdateExecutionContext(null, mockTestMetadata());
    UpdateEntityResponse response = entityUpdater.update(updateEntityRequest, updateExecutionContext).build();
    assertFalse(response.hasEntity());
}

Example 3 with EntityQueryServiceClient

use of org.hypertrace.entity.query.service.client.EntityQueryServiceClient in project gateway-service by hypertrace.

the class EdsEntityUpdaterTest method updateSuccess.

@Test
public void updateSuccess() {
    EntityQueryServiceClient mockEqsClient = mock(EntityQueryServiceClient.class);
    EdsEntityUpdater entityUpdater = new EdsEntityUpdater(mockEqsClient);
    // mock ResultSetChunk response from EQS
    UpdateEntityRequest updateEntityRequest = createResolveTestRequest("some-id");
    ResultSetChunk result = ResultSetChunk.newBuilder().addRow(Row.newBuilder().addColumn(org.hypertrace.entity.query.service.v1.Value.newBuilder().setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING).setString("some-id")).addColumn(org.hypertrace.entity.query.service.v1.Value.newBuilder().setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING).setString("RESOLVED"))).setResultSetMetadata(ResultSetMetadata.newBuilder().addColumnMetadata(ColumnMetadata.newBuilder().setColumnName("Test.id")).addColumnMetadata(ColumnMetadata.newBuilder().setColumnName("Test.status"))).build();
    when(mockEqsClient.update(any(), any())).thenReturn(List.of(result).iterator());
    UpdateExecutionContext updateExecutionContext = new UpdateExecutionContext(null, mockTestMetadata());
    UpdateEntityResponse response = entityUpdater.update(updateEntityRequest, updateExecutionContext).build();
    // verify response
    assertTrue(response.hasEntity());
    assertEquals(BACKEND.name(), response.getEntity().getEntityType());
    assertEquals(2, response.getEntity().getAttributeCount());
    assertEquals(ValueType.STRING, response.getEntity().getAttributeMap().get("Test.id").getValueType());
    assertEquals("some-id", response.getEntity().getAttributeMap().get("Test.id").getString());
    assertEquals(ValueType.STRING, response.getEntity().getAttributeMap().get("Test.status").getValueType());
    assertEquals("RESOLVED", response.getEntity().getAttributeMap().get("Test.status").getString());
    // verify request that goes to EQS
    EntityUpdateRequest expectedEqsUpdateRequest = EntityUpdateRequest.newBuilder().setEntityType(BACKEND.name()).addEntityIds("some-id").setOperation(UpdateOperation.newBuilder().setSetAttribute(org.hypertrace.entity.query.service.v1.SetAttribute.newBuilder().setAttribute(org.hypertrace.entity.query.service.v1.ColumnIdentifier.newBuilder().setColumnName("Test.status")).setValue(org.hypertrace.entity.query.service.v1.LiteralConstant.newBuilder().setValue(org.hypertrace.entity.query.service.v1.Value.newBuilder().setValueType(org.hypertrace.entity.query.service.v1.ValueType.STRING).setString("RESOLVED"))))).addSelection(org.hypertrace.entity.query.service.v1.Expression.newBuilder().setColumnIdentifier(org.hypertrace.entity.query.service.v1.ColumnIdentifier.newBuilder().setColumnName("Test.id"))).addSelection(org.hypertrace.entity.query.service.v1.Expression.newBuilder().setColumnIdentifier(org.hypertrace.entity.query.service.v1.ColumnIdentifier.newBuilder().setColumnName("Test.status"))).build();
    verify(mockEqsClient, times(1)).update(eq(expectedEqsUpdateRequest), any());
}