atlas-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From shweth...@apache.org
Subject [05/12] incubator-atlas git commit: ATLAS-947 Return state information in inputs and outputs lineage API (shwethags)
Date Wed, 29 Jun 2016 08:37:51 GMT
ATLAS-947 Return state information in inputs and outputs lineage API (shwethags)


Project: http://git-wip-us.apache.org/repos/asf/incubator-atlas/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-atlas/commit/40cccc36
Tree: http://git-wip-us.apache.org/repos/asf/incubator-atlas/tree/40cccc36
Diff: http://git-wip-us.apache.org/repos/asf/incubator-atlas/diff/40cccc36

Branch: refs/heads/0.7-incubating
Commit: 40cccc3662d5d78d35b128ecb475ea25c0c48855
Parents: 41ebb5e
Author: Shwetha GS <sshivalingamurthy@hortonworks.com>
Authored: Wed Jun 29 12:09:22 2016 +0530
Committer: Shwetha GS <sshivalingamurthy@hortonworks.com>
Committed: Wed Jun 29 12:09:22 2016 +0530

----------------------------------------------------------------------
 .../org/apache/atlas/hive/hook/HiveHookIT.java  |  3 +-
 release-log.txt                                 |  1 +
 .../atlas/discovery/DataSetLineageService.java  |  3 +-
 .../graph/DefaultGraphPersistenceStrategy.java  |  7 ++-
 .../atlas/repository/MetadataRepository.java    |  6 ++
 .../graph/GraphBackedMetadataRepository.java    |  5 ++
 .../query/GraphPersistenceStrategies.scala      |  6 ++
 .../discovery/DataSetLineageServiceTest.java    | 23 +++++--
 .../atlas/typesystem/types/TypeSystem.java      | 37 +++++++-----
 .../apache/atlas/ApplicationPropertiesTest.java |  4 +-
 .../test/resources/atlas-application.properties |  2 +-
 .../DataSetLineageJerseyResourceIT.java         | 63 ++++++++++----------
 12 files changed, 104 insertions(+), 56 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/addons/hive-bridge/src/test/java/org/apache/atlas/hive/hook/HiveHookIT.java
----------------------------------------------------------------------
diff --git a/addons/hive-bridge/src/test/java/org/apache/atlas/hive/hook/HiveHookIT.java b/addons/hive-bridge/src/test/java/org/apache/atlas/hive/hook/HiveHookIT.java
index 2464646..f9e1926 100755
--- a/addons/hive-bridge/src/test/java/org/apache/atlas/hive/hook/HiveHookIT.java
+++ b/addons/hive-bridge/src/test/java/org/apache/atlas/hive/hook/HiveHookIT.java
@@ -536,8 +536,7 @@ public class HiveHookIT {
         Referenceable processRef1 = validateProcess(query, HiveOperation.QUERY, inputs, outputs);
 
         //Rerun same query. Should result in same process
-        runCommand(query);
-
+        runCommandWithDelay(query, 1000);
         Referenceable processRef2 = validateProcess(query, HiveOperation.QUERY, inputs, outputs);
         Assert.assertEquals(processRef1.getId()._getId(), processRef2.getId()._getId());
 

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/release-log.txt
----------------------------------------------------------------------
diff --git a/release-log.txt b/release-log.txt
index 8ccbac9..b4e8fc5 100644
--- a/release-log.txt
+++ b/release-log.txt
@@ -31,6 +31,7 @@ ATLAS-409 Atlas will not import avro tables with schema read from a file
(dosset
 ATLAS-379 Create sqoop and falcon metadata addons (venkatnrangan,bvellanki,sowmyaramesh via
shwethags)
 
 ALL CHANGES:
+ATLAS-947 Return state information in inputs and outputs lineage API (shwethags)
 ATLAS-806 Create default taxonomy at server startup (jspeidel via yhemanth)
 ATLAS-942 Jenkins build failure - GraphRepoMapperScaleTest (shwethags)
 ATLAS-920 Lineage graph is broken when there are multiple paths from same source table (kevalbhatt18
via sumasai)

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/repository/src/main/java/org/apache/atlas/discovery/DataSetLineageService.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/discovery/DataSetLineageService.java
b/repository/src/main/java/org/apache/atlas/discovery/DataSetLineageService.java
index a2f47f0..c6790de 100644
--- a/repository/src/main/java/org/apache/atlas/discovery/DataSetLineageService.java
+++ b/repository/src/main/java/org/apache/atlas/discovery/DataSetLineageService.java
@@ -52,7 +52,8 @@ public class DataSetLineageService implements LineageService {
     private static final Logger LOG = LoggerFactory.getLogger(DataSetLineageService.class);
 
     private static final Option<List<String>> SELECT_ATTRIBUTES =
-            Some.<List<String>>apply(List.<String>fromArray(new String[]{"name"}));
+            Some.<List<String>>apply(List.<String>fromArray(new String[]{AtlasClient.NAME,
+                    AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME}));
     public static final String SELECT_INSTANCE_GUID = "__guid";
 
     public static final String DATASET_SCHEMA_QUERY_PREFIX = "atlas.lineage.schema.query.";

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java
b/repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java
index 90718ed..8ddbe1b 100755
--- a/repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java
+++ b/repository/src/main/java/org/apache/atlas/discovery/graph/DefaultGraphPersistenceStrategy.java
@@ -142,7 +142,7 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
                 if (dataType.getName().equals(idType.getName())) {
                     structInstance.set(idType.typeNameAttrName(), structVertex.getProperty(typeAttributeName()));
                     structInstance.set(idType.idAttrName(), structVertex.getProperty(idAttributeName()));
-
+                    structInstance.set(idType.stateAttrName(), structVertex.getProperty(stateAttributeName()));
                 } else {
                     metadataRepository.getGraphToInstanceMapper()
                         .mapVertexToInstance(structVertex, structInstance, structType.fieldMapping().fields);
@@ -229,6 +229,11 @@ public class DefaultGraphPersistenceStrategy implements GraphPersistenceStrategi
     }
 
     @Override
+    public String stateAttributeName() {
+        return metadataRepository.getStateAttributeName();
+    }
+
+    @Override
     public scala.collection.Seq<String> typeTestExpression(String typeName, IntSequence
intSeq) {
         return GraphPersistenceStrategies$class.typeTestExpression(this, typeName, intSeq);
     }

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java
b/repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java
index 43e9f85..a2a9ab6 100755
--- a/repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java
+++ b/repository/src/main/java/org/apache/atlas/repository/MetadataRepository.java
@@ -50,6 +50,12 @@ public interface MetadataRepository {
     String getSuperTypeAttributeName();
 
     /**
+     * Returns the attribute name used for entity state
+     * @return
+     */
+    String getStateAttributeName();
+
+    /**
      * Return the property key used to store a given traitName in the repository.
      *
      * @param dataType  data type

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java
b/repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java
index a4bdef7..e5c0e46 100755
--- a/repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java
+++ b/repository/src/main/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepository.java
@@ -86,6 +86,11 @@ public class GraphBackedMetadataRepository implements MetadataRepository
{
         return Constants.ENTITY_TYPE_PROPERTY_KEY;
     }
 
+    @Override
+    public String getStateAttributeName() {
+        return Constants.STATE_PROPERTY_KEY;
+    }
+
     /**
      * Returns the property key used to store super type names.
      *

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/repository/src/main/scala/org/apache/atlas/query/GraphPersistenceStrategies.scala
----------------------------------------------------------------------
diff --git a/repository/src/main/scala/org/apache/atlas/query/GraphPersistenceStrategies.scala
b/repository/src/main/scala/org/apache/atlas/query/GraphPersistenceStrategies.scala
index 34d101a..21c152c 100755
--- a/repository/src/main/scala/org/apache/atlas/query/GraphPersistenceStrategies.scala
+++ b/repository/src/main/scala/org/apache/atlas/query/GraphPersistenceStrategies.scala
@@ -62,6 +62,11 @@ trait GraphPersistenceStrategies {
     def idAttributeName : String
 
     /**
+      * Name of attribute used to store state in vertex
+      */
+    def stateAttributeName : String
+
+    /**
      * Given a dataType and a reference attribute, how is edge labeled
      */
     def edgeLabel(iDataType: IDataType[_], aInfo: AttributeInfo): String
@@ -190,6 +195,7 @@ object GraphPersistenceStrategy1 extends GraphPersistenceStrategies {
     val typeAttributeName = "typeName"
     val superTypeAttributeName = "superTypeNames"
     val idAttributeName = "guid"
+    val stateAttributeName = "state"
 
     def edgeLabel(dataType: IDataType[_], aInfo: AttributeInfo) = s"__${dataType.getName}.${aInfo.name}"
 

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
b/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
index 3028dde..89e9b41 100644
--- a/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
+++ b/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
@@ -19,12 +19,15 @@
 package org.apache.atlas.discovery;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.BaseRepositoryTest;
 import org.apache.atlas.RepositoryMetadataModule;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.Struct;
 import org.apache.atlas.typesystem.exception.EntityNotFoundException;
+import org.apache.atlas.typesystem.json.InstanceSerialization;
 import org.apache.atlas.typesystem.persistence.Id;
 import org.apache.commons.collections.ArrayStack;
 import org.apache.commons.lang.RandomStringUtils;
@@ -40,9 +43,11 @@ import org.testng.annotations.Test;
 import javax.inject.Inject;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 
 /**
@@ -334,18 +339,21 @@ public class DataSetLineageServiceTest extends BaseRepositoryTest {
     public void testLineageWithDelete() throws Exception {
         String tableName = "table" + random();
         createTable(tableName, 3, true);
+        String tableId = getEntityId(HIVE_TABLE_TYPE, "name", tableName);
 
         JSONObject results = new JSONObject(lineageService.getSchema(tableName));
         assertEquals(results.getJSONArray("rows").length(), 3);
 
         results = new JSONObject(lineageService.getInputsGraph(tableName));
-        assertEquals(results.getJSONObject("values").getJSONObject("vertices").length(),
2);
+        Struct resultInstance = InstanceSerialization.fromJsonStruct(results.toString(),
true);
+        Map<String, Struct> vertices = (Map) resultInstance.get("vertices");
+        assertEquals(vertices.size(), 2);
+        Struct vertex = vertices.get(tableId);
+        assertEquals(((Struct) vertex.get("vertexId")).get("state"), Id.EntityState.ACTIVE.name());
 
         results = new JSONObject(lineageService.getOutputsGraph(tableName));
         assertEquals(results.getJSONObject("values").getJSONObject("vertices").length(),
2);
 
-        String tableId = getEntityId(HIVE_TABLE_TYPE, "name", tableName);
-
         results = new JSONObject(lineageService.getSchemaForEntity(tableId));
         assertEquals(results.getJSONArray("rows").length(), 3);
 
@@ -357,12 +365,19 @@ public class DataSetLineageServiceTest extends BaseRepositoryTest {
 
         //Delete the entity. Lineage for entity returns the same results as before.
         //Lineage for table name throws EntityNotFoundException
-        repository.deleteEntities(Arrays.asList(tableId));
+        AtlasClient.EntityResult deleteResult = repository.deleteEntities(Arrays.asList(tableId));
+        assertTrue(deleteResult.getDeletedEntities().contains(tableId));
 
         results = new JSONObject(lineageService.getSchemaForEntity(tableId));
         assertEquals(results.getJSONArray("rows").length(), 3);
 
         results = new JSONObject(lineageService.getInputsGraphForEntity(tableId));
+        resultInstance = InstanceSerialization.fromJsonStruct(results.toString(), true);
+        vertices = (Map) resultInstance.get("vertices");
+        assertEquals(vertices.size(), 2);
+        vertex = vertices.get(tableId);
+        assertEquals(((Struct) vertex.get("vertexId")).get("state"), Id.EntityState.DELETED.name());
+
         assertEquals(results.getJSONObject("values").getJSONObject("vertices").length(),
2);
 
         results = new JSONObject(lineageService.getOutputsGraphForEntity(tableId));

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeSystem.java
----------------------------------------------------------------------
diff --git a/typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeSystem.java b/typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeSystem.java
index 90b5815..c355ba9 100755
--- a/typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeSystem.java
+++ b/typesystem/src/main/java/org/apache/atlas/typesystem/types/TypeSystem.java
@@ -18,18 +18,8 @@
 
 package org.apache.atlas.typesystem.types;
 
-import java.lang.reflect.Constructor;
-import java.text.SimpleDateFormat;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.TimeZone;
-import java.util.concurrent.ConcurrentHashMap;
-
-import javax.inject.Singleton;
-
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.classification.InterfaceAudience;
 import org.apache.atlas.typesystem.TypesDef;
@@ -40,8 +30,16 @@ import org.apache.atlas.typesystem.types.cache.TypeCache;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
+import javax.inject.Singleton;
+import java.lang.reflect.Constructor;
+import java.text.SimpleDateFormat;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TimeZone;
+import java.util.concurrent.ConcurrentHashMap;
 
 @Singleton
 @InterfaceAudience.Private
@@ -717,6 +715,7 @@ public class TypeSystem {
     public class IdType {
         private static final String ID_ATTRNAME = "guid";
         private static final String TYPENAME_ATTRNAME = "typeName";
+        private static final String STATE_ATTRNAME = "state";
         private static final String TYP_NAME = "__IdType";
 
         private StructType type;
@@ -728,10 +727,14 @@ public class TypeSystem {
             AttributeDefinition typNmAttr =
                     new AttributeDefinition(TYPENAME_ATTRNAME, DataTypes.STRING_TYPE.getName(),
Multiplicity.REQUIRED,
                             false, null);
+            AttributeDefinition stateAttr =
+                    new AttributeDefinition(STATE_ATTRNAME, DataTypes.STRING_TYPE.getName(),
Multiplicity.REQUIRED,
+                            false, null);
             try {
-                AttributeInfo[] infos = new AttributeInfo[2];
+                AttributeInfo[] infos = new AttributeInfo[3];
                 infos[0] = new AttributeInfo(TypeSystem.this, idAttr, null);
                 infos[1] = new AttributeInfo(TypeSystem.this, typNmAttr, null);
+                infos[2] = new AttributeInfo(TypeSystem.this, stateAttr, null);
 
                 type = new StructType(TypeSystem.this, TYP_NAME, null, infos);
             } catch (AtlasException me) {
@@ -754,6 +757,10 @@ public class TypeSystem {
         public String typeNameAttrName() {
             return TYPENAME_ATTRNAME;
         }
+
+        public String stateAttrName() {
+            return STATE_ATTRNAME;
+        }
     }
 
     public static final String ID_STRUCT_ID_ATTRNAME = IdType.ID_ATTRNAME;

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/typesystem/src/test/java/org/apache/atlas/ApplicationPropertiesTest.java
----------------------------------------------------------------------
diff --git a/typesystem/src/test/java/org/apache/atlas/ApplicationPropertiesTest.java b/typesystem/src/test/java/org/apache/atlas/ApplicationPropertiesTest.java
index 6d43359..25a19cf 100644
--- a/typesystem/src/test/java/org/apache/atlas/ApplicationPropertiesTest.java
+++ b/typesystem/src/test/java/org/apache/atlas/ApplicationPropertiesTest.java
@@ -34,7 +34,7 @@ public class ApplicationPropertiesTest {
         assertEquals(properties.getString("atlas.service"), "atlas");
 
         //property containing system property
-        String data = "/var/data/" + System.getProperty("user.name") + "/atlas";
+        String data = System.getProperty("user.dir") + "/target/data";
         assertEquals(properties.getString("atlas.data"), data);
 
         //property referencing other property
@@ -51,7 +51,7 @@ public class ApplicationPropertiesTest {
         Configuration subConfiguration = configuration.subset("atlas");
 
         assertEquals(subConfiguration.getString("service"), "atlas");
-        String data = "/var/data/" + System.getProperty("user.name") + "/atlas";
+        String data = System.getProperty("user.dir") + "/target/data";
         assertEquals(subConfiguration.getString("data"), data);
         assertEquals(subConfiguration.getString("graph.data"), data + "/graph");
     }

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/typesystem/src/test/resources/atlas-application.properties
----------------------------------------------------------------------
diff --git a/typesystem/src/test/resources/atlas-application.properties b/typesystem/src/test/resources/atlas-application.properties
index 28d8af3..4d7e9b5 100644
--- a/typesystem/src/test/resources/atlas-application.properties
+++ b/typesystem/src/test/resources/atlas-application.properties
@@ -17,7 +17,7 @@
 #
 
 #system property
-atlas.data=/var/data/${sys:user.name}/atlas
+atlas.data=${sys:user.dir}/target/data
 
 #re-use existing property
 atlas.graph.data=${atlas.data}/graph

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/40cccc36/webapp/src/test/java/org/apache/atlas/web/resources/DataSetLineageJerseyResourceIT.java
----------------------------------------------------------------------
diff --git a/webapp/src/test/java/org/apache/atlas/web/resources/DataSetLineageJerseyResourceIT.java
b/webapp/src/test/java/org/apache/atlas/web/resources/DataSetLineageJerseyResourceIT.java
index bc4f239..0ffaa1f 100644
--- a/webapp/src/test/java/org/apache/atlas/web/resources/DataSetLineageJerseyResourceIT.java
+++ b/webapp/src/test/java/org/apache/atlas/web/resources/DataSetLineageJerseyResourceIT.java
@@ -23,6 +23,8 @@ import com.sun.jersey.api.client.ClientResponse;
 import com.sun.jersey.api.client.WebResource;
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.typesystem.Referenceable;
+import org.apache.atlas.typesystem.Struct;
+import org.apache.atlas.typesystem.json.InstanceSerialization;
 import org.apache.atlas.typesystem.persistence.Id;
 import org.apache.atlas.web.util.Servlets;
 import org.codehaus.jettison.json.JSONArray;
@@ -34,6 +36,9 @@ import org.testng.annotations.Test;
 import javax.ws.rs.HttpMethod;
 import javax.ws.rs.core.Response;
 import java.util.List;
+import java.util.Map;
+
+import static org.testng.Assert.assertEquals;
 
 /**
  * Hive Lineage Integration Tests.
@@ -70,30 +75,29 @@ public class DataSetLineageJerseyResourceIT extends BaseResourceIT {
         JSONObject results = response.getJSONObject(AtlasClient.RESULTS);
         Assert.assertNotNull(results);
 
-        JSONObject values = results.getJSONObject("values");
-        Assert.assertNotNull(values);
-
-        final JSONObject vertices = values.getJSONObject("vertices");
-        Assert.assertEquals(vertices.length(), 4);
+        Struct resultsInstance = InstanceSerialization.fromJsonStruct(results.toString(),
true);
+        Map<String, Struct> vertices = (Map<String, Struct>) resultsInstance.get("vertices");
+        Assert.assertEquals(vertices.size(), 4);
 
-        final JSONObject edges = values.getJSONObject("edges");
-        Assert.assertEquals(edges.length(), 4);
+        Map<String, Struct> edges = (Map<String, Struct>) resultsInstance.get("edges");
+        Assert.assertEquals(edges.size(), 4);
     }
 
     @Test
     public void testInputsGraphForEntity() throws Exception {
-        String tableId = serviceClient.getEntity(HIVE_TABLE_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
salesMonthlyTable).getId()._getId();
+        String tableId = serviceClient.getEntity(HIVE_TABLE_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
+                salesMonthlyTable).getId()._getId();
         JSONObject results = serviceClient.getInputGraphForEntity(tableId);
         Assert.assertNotNull(results);
 
-        JSONObject values = results.getJSONObject("values");
-        Assert.assertNotNull(values);
-
-        final JSONObject vertices = values.getJSONObject("vertices");
-        Assert.assertEquals(vertices.length(), 4);
+        Struct resultsInstance = InstanceSerialization.fromJsonStruct(results.toString(),
true);
+        Map<String, Struct> vertices = (Map<String, Struct>) resultsInstance.get("vertices");
+        Assert.assertEquals(vertices.size(), 4);
+        Struct vertex = vertices.get(tableId);
+        assertEquals(((Struct) vertex.get("vertexId")).get("state"), Id.EntityState.ACTIVE.name());
 
-        final JSONObject edges = values.getJSONObject("edges");
-        Assert.assertEquals(edges.length(), 4);
+        Map<String, Struct> edges = (Map<String, Struct>) resultsInstance.get("edges");
+        Assert.assertEquals(edges.size(), 4);
     }
 
     @Test
@@ -114,30 +118,29 @@ public class DataSetLineageJerseyResourceIT extends BaseResourceIT {
         JSONObject results = response.getJSONObject(AtlasClient.RESULTS);
         Assert.assertNotNull(results);
 
-        JSONObject values = results.getJSONObject("values");
-        Assert.assertNotNull(values);
+        Struct resultsInstance = InstanceSerialization.fromJsonStruct(results.toString(),
true);
+        Map<String, Struct> vertices = (Map<String, Struct>) resultsInstance.get("vertices");
+        Assert.assertEquals(vertices.size(), 3);
 
-        final JSONObject vertices = values.getJSONObject("vertices");
-        Assert.assertEquals(vertices.length(), 3);
-
-        final JSONObject edges = values.getJSONObject("edges");
-        Assert.assertEquals(edges.length(), 4);
+        Map<String, Struct> edges = (Map<String, Struct>) resultsInstance.get("edges");
+        Assert.assertEquals(edges.size(), 4);
     }
 
     @Test
     public void testOutputsGraphForEntity() throws Exception {
-        String tableId = serviceClient.getEntity(HIVE_TABLE_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
salesFactTable).getId()._getId();
+        String tableId = serviceClient.getEntity(HIVE_TABLE_TYPE, AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME,
+                salesFactTable).getId()._getId();
         JSONObject results = serviceClient.getOutputGraphForEntity(tableId);
         Assert.assertNotNull(results);
 
-        JSONObject values = results.getJSONObject("values");
-        Assert.assertNotNull(values);
-
-        final JSONObject vertices = values.getJSONObject("vertices");
-        Assert.assertEquals(vertices.length(), 3);
+        Struct resultsInstance = InstanceSerialization.fromJsonStruct(results.toString(),
true);
+        Map<String, Struct> vertices = (Map<String, Struct>) resultsInstance.get("vertices");
+        Assert.assertEquals(vertices.size(), 3);
+        Struct vertex = vertices.get(tableId);
+        assertEquals(((Struct) vertex.get("vertexId")).get("state"), Id.EntityState.ACTIVE.name());
 
-        final JSONObject edges = values.getJSONObject("edges");
-        Assert.assertEquals(edges.length(), 4);
+        Map<String, Struct> edges = (Map<String, Struct>) resultsInstance.get("edges");
+        Assert.assertEquals(edges.size(), 4);
     }
 
     @Test


Mime
View raw message