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));
  }