private String getSample(AtlasClient.EntityResult entityResult) {
String sample = getSample(entityResult.getCreatedEntities());
if (sample == null) {
sample = getSample(entityResult.getUpdateEntities());
}
return sample;
}
@Test
public void testUpdateEntity() throws Exception {
final String guid = random();
Response response = mock(Response.class);
when(entityResource.updateByUniqueAttribute(
anyString(), anyString(), anyString(), any(HttpServletRequest.class)))
.thenReturn(response);
when(response.getEntity())
.thenReturn(
new JSONObject() {
{
put(
ENTITIES,
new JSONObject(
new AtlasClient.EntityResult(null, Arrays.asList(guid), null)
.toString())
.get(ENTITIES));
}
});
LocalAtlasClient atlasClient = new LocalAtlasClient(serviceState, entityResource);
AtlasClient.EntityResult entityResult =
atlasClient.updateEntity(random(), random(), random(), new Referenceable(random()));
assertEquals(entityResult.getUpdateEntities(), Arrays.asList(guid));
}
private JSONObject getResponse(AtlasClient.EntityResult entityResult)
throws AtlasException, JSONException {
JSONObject response = new JSONObject();
response.put(AtlasClient.REQUEST_ID, Servlets.getRequestId());
response.put(
AtlasClient.ENTITIES, new JSONObject(entityResult.toString()).get(AtlasClient.ENTITIES));
String sampleEntityId = getSample(entityResult);
if (sampleEntityId != null) {
String entityDefinition = metadataService.getEntityDefinition(sampleEntityId);
response.put(AtlasClient.DEFINITION, new JSONObject(entityDefinition));
}
return response;
}
@Test
public void testDeleteEntity() throws Exception {
final String guid = random();
Response response = mock(Response.class);
when(response.getEntity())
.thenReturn(
new JSONObject() {
{
put(
ENTITIES,
new JSONObject(
new AtlasClient.EntityResult(null, null, Arrays.asList(guid))
.toString())
.get(ENTITIES));
}
});
when(entityResource.deleteEntities(
anyListOf(String.class), anyString(), anyString(), anyString()))
.thenReturn(response);
LocalAtlasClient atlasClient = new LocalAtlasClient(serviceState, entityResource);
AtlasClient.EntityResult entityResult = atlasClient.deleteEntity(random(), random(), random());
assertEquals(entityResult.getDeletedEntities(), Arrays.asList(guid));
}