You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@atlas.apache.org by ap...@apache.org on 2017/05/23 22:07:08 UTC

[10/12] incubator-atlas git commit: ATLAS-1198: Spring Framework (v4 with Spring security) over Guice

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java
index 70a904b..27c0b5d 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStoreV1.java
@@ -18,11 +18,9 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 
-import com.google.inject.Inject;
-import com.google.inject.Singleton;
 import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.GraphTransaction;
 import org.apache.atlas.RequestContextV1;
+import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.impexp.AtlasImportResult;
 import org.apache.atlas.model.instance.AtlasClassification;
@@ -47,7 +45,9 @@ import org.apache.commons.collections.MapUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
 
+import javax.inject.Inject;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -61,19 +61,22 @@ import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DE
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE;
 
 
-@Singleton
+@Component
 public class AtlasEntityStoreV1 implements AtlasEntityStore {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV1.class);
 
     private final DeleteHandlerV1           deleteHandler;
     private final AtlasTypeRegistry         typeRegistry;
     private final AtlasEntityChangeNotifier entityChangeNotifier;
+    private final EntityGraphMapper entityGraphMapper;
 
     @Inject
-    public AtlasEntityStoreV1(DeleteHandlerV1 deleteHandler, AtlasTypeRegistry typeRegistry, AtlasEntityChangeNotifier entityChangeNotifier) {
+    public AtlasEntityStoreV1(DeleteHandlerV1 deleteHandler, AtlasTypeRegistry typeRegistry,
+                              AtlasEntityChangeNotifier entityChangeNotifier, EntityGraphMapper entityGraphMapper) {
         this.deleteHandler        = deleteHandler;
         this.typeRegistry         = typeRegistry;
         this.entityChangeNotifier = entityChangeNotifier;
+        this.entityGraphMapper = entityGraphMapper;
     }
 
     @Override
@@ -143,6 +146,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
+    @GraphTransaction
     public EntityMutationResponse bulkImport(EntityImportStream entityStream, AtlasImportResult importResult) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> bulkImport()");
@@ -208,7 +212,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
         }
     }
 
-    @GraphTransaction
     private EntityMutationResponse createOrUpdate(EntityStream entityStream, boolean isPartialUpdate, boolean replaceClassifications) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> createOrUpdate()");
@@ -219,8 +222,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
         }
 
         // Create/Update entities
-        EntityGraphMapper entityGraphMapper = new EntityGraphMapper(deleteHandler, typeRegistry);
-
         EntityMutationContext context = preCreateOrUpdate(entityStream, entityGraphMapper, isPartialUpdate);
 
         EntityMutationResponse ret = entityGraphMapper.mapAttributesAndClassifications(context, isPartialUpdate, replaceClassifications);
@@ -238,6 +239,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
     }
 
     @Override
+    @GraphTransaction
     public EntityMutationResponse createOrUpdate(EntityStream entityStream, boolean isPartialUpdate) throws AtlasBaseException {
         return createOrUpdate(entityStream, isPartialUpdate, false);
     }
@@ -315,6 +317,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
         return createOrUpdate(new AtlasEntityStream(updateEntity), true);
     }
 
+    @Override
     @GraphTransaction
     public EntityMutationResponse deleteById(final String guid) throws AtlasBaseException {
 
@@ -432,8 +435,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
         // validate if entity, not already associated with classifications
         validateEntityAssociations(guid, classifications);
 
-        EntityGraphMapper graphMapper = new EntityGraphMapper(deleteHandler, typeRegistry);
-        graphMapper.addClassifications(new EntityMutationContext(), guid, classifications);
+        entityGraphMapper.addClassifications(new EntityMutationContext(), guid, classifications);
 
         // notify listeners on classification addition
         entityChangeNotifier.onClassificationAddedToEntity(guid, classifications);
@@ -454,7 +456,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "classifications(s) not specified");
         }
 
-        EntityGraphMapper         graphMapper            = new EntityGraphMapper(deleteHandler, typeRegistry);
         List<AtlasClassification> updatedClassifications = new ArrayList<>();
 
         for (AtlasClassification newClassification : newClassifications) {
@@ -475,7 +476,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
                 }
             }
 
-            graphMapper.updateClassification(new EntityMutationContext(), guid, oldClassification);
+            entityGraphMapper.updateClassification(new EntityMutationContext(), guid, oldClassification);
 
             updatedClassifications.add(oldClassification);
         }
@@ -498,8 +499,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             LOG.debug("Adding classification={} to entities={}", classification, guids);
         }
 
-        EntityGraphMapper graphMapper = new EntityGraphMapper(deleteHandler, typeRegistry);
-
         validateAndNormalize(classification);
 
         List<AtlasClassification> classifications = Collections.singletonList(classification);
@@ -508,7 +507,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             // validate if entity, not already associated with classifications
             validateEntityAssociations(guid, classifications);
 
-            graphMapper.addClassifications(new EntityMutationContext(), guid, classifications);
+            entityGraphMapper.addClassifications(new EntityMutationContext(), guid, classifications);
 
             // notify listeners on classification addition
             entityChangeNotifier.onClassificationAddedToEntity(guid, classifications);
@@ -530,7 +529,6 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             LOG.debug("Deleting classifications={} from entity={}", classificationNames, guid);
         }
 
-        EntityGraphMapper entityGraphMapper = new EntityGraphMapper(deleteHandler, typeRegistry);
         entityGraphMapper.deleteClassifications(guid, classificationNames);
 
         // notify listeners on classification deletion
@@ -569,7 +567,7 @@ public class AtlasEntityStoreV1 implements AtlasEntityStore {
             AtlasEntity entity = entityStream.getByGuid(guid);
 
             if (entity != null) {
-                
+
                 if (vertex != null) {
                     // entity would be null if guid is not in the stream but referenced by an entity in the stream
                     if (!isPartialUpdate) {

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStream.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStream.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStream.java
index eb860ff..d6d6272 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStream.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasEntityStream.java
@@ -18,8 +18,8 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 import org.apache.atlas.model.instance.AtlasEntity;
-import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
+import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 
 import java.util.Iterator;
 

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasTypeDefGraphStoreV1.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasTypeDefGraphStoreV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasTypeDefGraphStoreV1.java
index 3bf318f..f0c8380 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasTypeDefGraphStoreV1.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/AtlasTypeDefGraphStoreV1.java
@@ -19,14 +19,11 @@ package org.apache.atlas.repository.store.graph.v1;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
-import com.google.inject.Inject;
-
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.listener.TypeDefChangeListener;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.repository.Constants;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
@@ -43,7 +40,11 @@ import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
 
+import javax.annotation.PostConstruct;
+import javax.inject.Inject;
+import javax.inject.Singleton;
 import java.util.Date;
 import java.util.HashSet;
 import java.util.Iterator;
@@ -59,15 +60,19 @@ import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.VERTE
 /**
  * Graph persistence store for TypeDef - v1
  */
+@Singleton
+@Component
 public class AtlasTypeDefGraphStoreV1 extends AtlasTypeDefGraphStore {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasTypeDefGraphStoreV1.class);
 
-    protected final AtlasGraph atlasGraph = AtlasGraphProvider.getGraphInstance();
+    protected final AtlasGraph atlasGraph;
 
     @Inject
     public AtlasTypeDefGraphStoreV1(AtlasTypeRegistry typeRegistry,
-                                    Set<TypeDefChangeListener> typeDefChangeListeners) {
+                                    Set<TypeDefChangeListener> typeDefChangeListeners,
+                                    AtlasGraph atlasGraph) {
         super(typeRegistry, typeDefChangeListeners);
+        this.atlasGraph = atlasGraph;
 
         LOG.debug("==> AtlasTypeDefGraphStoreV1()");
 
@@ -105,6 +110,7 @@ public class AtlasTypeDefGraphStoreV1 extends AtlasTypeDefGraphStore {
     }
 
     @Override
+    @PostConstruct
     public void init() throws AtlasBaseException {
         LOG.debug("==> AtlasTypeDefGraphStoreV1.init()");
 

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
index 0fe748e..80cd1ee 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphMapper.java
@@ -18,7 +18,6 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 
-import com.google.inject.Inject;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.RequestContextV1;
@@ -29,9 +28,8 @@ import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
-import org.apache.atlas.model.instance.EntityMutations;
-import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
 import org.apache.atlas.model.instance.EntityMutationResponse;
+import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.RepositoryException;
@@ -45,48 +43,42 @@ import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasMapType;
 import org.apache.atlas.type.AtlasStructType;
+import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.type.AtlasTypeUtil;
-import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.UUID;
+import javax.inject.Inject;
+import java.util.*;
 
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.CREATE;
+import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DELETE;
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.PARTIAL_UPDATE;
 import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.UPDATE;
-import static org.apache.atlas.model.instance.EntityMutations.EntityOperation.DELETE;
 import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 import static org.apache.atlas.repository.graph.GraphHelper.string;
 
-
+@Component
 public class EntityGraphMapper {
     private static final Logger LOG = LoggerFactory.getLogger(EntityGraphMapper.class);
 
-    private final AtlasGraph        graph       = AtlasGraphProvider.getGraphInstance();
     private final GraphHelper       graphHelper = GraphHelper.getInstance();
+    private final AtlasGraph        graph;
     private final DeleteHandlerV1   deleteHandler;
     private final AtlasTypeRegistry typeRegistry;
 
 
     @Inject
-    public EntityGraphMapper(DeleteHandlerV1 deleteHandler, AtlasTypeRegistry typeRegistry) {
+    public EntityGraphMapper(DeleteHandlerV1 deleteHandler, AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph) {
         this.deleteHandler = deleteHandler;
         this.typeRegistry  = typeRegistry;
+        this.graph         = atlasGraph;
     }
 
     public AtlasVertex createVertex(AtlasEntity entity) {
@@ -284,7 +276,7 @@ public class EntityGraphMapper {
     }
 
     private Object mapToVertexByTypeCategory(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
-        if (ctx.getOp() == EntityMutations.EntityOperation.CREATE && ctx.getValue() == null) {
+        if (ctx.getOp() == CREATE && ctx.getValue() == null) {
             return null;
         }
 
@@ -509,7 +501,7 @@ public class EntityGraphMapper {
 
         try {
             AtlasAttribute      attribute   = ctx.getAttribute();
-            List<String>        currentKeys = GraphHelper.getListProperty(ctx.getReferringVertex(), ctx.getVertexProperty());
+            List<String> currentKeys = GraphHelper.getListProperty(ctx.getReferringVertex(), ctx.getVertexProperty());
             Map<String, Object> currentMap  = new HashMap<>();
 
             if (CollectionUtils.isNotEmpty(currentKeys)) {

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java
index da6d636..66f20da 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/EntityGraphRetriever.java
@@ -289,16 +289,16 @@ public final class EntityGraphRetriever {
         List<String> classificationNames = GraphHelper.getTraitNames(instanceVertex);
 
         if (CollectionUtils.isNotEmpty(classificationNames)) {
-            for (String classficationName : classificationNames) {
-                AtlasClassification classification = null;
+            for (String classificationName : classificationNames) {
+                AtlasClassification classification;
                 if (StringUtils.isNotEmpty(classificationNameFilter)) {
-                    if (classficationName.equals(classificationNameFilter)) {
-                        classification = getClassification(instanceVertex, classficationName);
+                    if (classificationName.equals(classificationNameFilter)) {
+                        classification = getClassification(instanceVertex, classificationName);
                         classifications.add(classification);
                         return classifications;
                     }
                 } else {
-                    classification = getClassification(instanceVertex, classficationName);
+                    classification = getClassification(instanceVertex, classificationName);
                     classifications.add(classification);
                 }
             }

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java
index f51238a..29518c4 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java
@@ -18,12 +18,17 @@
 
 package org.apache.atlas.repository.store.graph.v1;
 
-import com.google.inject.Inject;
+import org.apache.atlas.annotation.ConditionalOnAtlasProperty;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasTypeRegistry;
+import org.springframework.stereotype.Component;
 
+import javax.inject.Inject;
+
+@Component
+@ConditionalOnAtlasProperty(property = "atlas.DeleteHandlerV1.impl")
 public class HardDeleteHandlerV1 extends DeleteHandlerV1 {
 
     @Inject

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java
index 31bb24a..49c9e0c 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java
@@ -18,21 +18,24 @@
 
 package org.apache.atlas.repository.store.graph.v1;
 
-import com.google.inject.Inject;
 import org.apache.atlas.RequestContextV1;
+import org.apache.atlas.annotation.ConditionalOnAtlasProperty;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
-import org.apache.atlas.RequestContext;
 import org.apache.atlas.type.AtlasTypeRegistry;
-import org.apache.atlas.typesystem.persistence.Id;
+import org.springframework.stereotype.Component;
+
+import javax.inject.Inject;
 
 import static org.apache.atlas.repository.Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY;
 import static org.apache.atlas.repository.Constants.MODIFIED_BY_KEY;
 import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 
+@Component
+@ConditionalOnAtlasProperty(property = "atlas.DeleteHandlerV1.impl", isDefault = true)
 public class SoftDeleteHandlerV1 extends DeleteHandlerV1 {
 
     @Inject

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/UniqAttrBasedEntityResolver.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/UniqAttrBasedEntityResolver.java b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/UniqAttrBasedEntityResolver.java
index c80639a..50eee72 100644
--- a/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/UniqAttrBasedEntityResolver.java
+++ b/repository/src/main/java/org/apache/atlas/repository/store/graph/v1/UniqAttrBasedEntityResolver.java
@@ -20,31 +20,22 @@ package org.apache.atlas.repository.store.graph.v1;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
-import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
-import org.apache.atlas.repository.Constants;
-import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.type.AtlasTypeRegistry;
-import org.apache.atlas.typesystem.exception.EntityNotFoundException;
-import org.apache.commons.collections.MapUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 
 
 public class UniqAttrBasedEntityResolver implements EntityResolver {
     private static final Logger LOG = LoggerFactory.getLogger(UniqAttrBasedEntityResolver.class);
 
-    private final static GraphHelper graphHelper = GraphHelper.getInstance();
-
     private final AtlasTypeRegistry typeRegistry;
 
     public UniqAttrBasedEntityResolver(AtlasTypeRegistry typeRegistry) {

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/typestore/GraphBackedTypeStore.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/typestore/GraphBackedTypeStore.java b/repository/src/main/java/org/apache/atlas/repository/typestore/GraphBackedTypeStore.java
index ac13586..7a064b6 100644
--- a/repository/src/main/java/org/apache/atlas/repository/typestore/GraphBackedTypeStore.java
+++ b/repository/src/main/java/org/apache/atlas/repository/typestore/GraphBackedTypeStore.java
@@ -22,13 +22,10 @@ import com.google.common.base.Function;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
-import com.google.inject.Inject;
-import com.google.inject.Singleton;
 import org.apache.atlas.AtlasException;
-import org.apache.atlas.GraphTransaction;
+import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.RepositoryException;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
@@ -41,7 +38,10 @@ import org.apache.atlas.typesystem.types.utils.TypesUtil;
 import org.codehaus.jettison.json.JSONException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
 
+import javax.inject.Inject;
+import javax.inject.Singleton;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -53,6 +53,7 @@ import java.util.Set;
 import static org.apache.atlas.repository.graph.GraphHelper.setProperty;
 
 @Singleton
+@Component
 @Deprecated
 public class GraphBackedTypeStore implements ITypeStore {
     public static final String VERTEX_TYPE = "typeSystem";
@@ -66,8 +67,8 @@ public class GraphBackedTypeStore implements ITypeStore {
     private GraphHelper graphHelper = GraphHelper.getInstance();
 
     @Inject
-    public GraphBackedTypeStore() {
-        graph = AtlasGraphProvider.getGraphInstance();
+    public GraphBackedTypeStore(AtlasGraph atlasGraph) {
+        this.graph = atlasGraph;
     }
 
     @Override
@@ -123,7 +124,7 @@ public class GraphBackedTypeStore implements ITypeStore {
         List<AtlasVertex> vertices = createVertices(typeVerticesNeeded);
 
         //Create a type name->AtlasVertex map with the result
-        Map<String, AtlasVertex> result = new HashMap<String,AtlasVertex>(typeVerticesNeeded.size());
+        Map<String, AtlasVertex> result = new HashMap<>(typeVerticesNeeded.size());
         for(int i = 0 ; i < typeVerticesNeeded.size(); i++) {
             TypeVertexInfo createdVertexInfo = typeVerticesNeeded.get(i);
             AtlasVertex createdVertex = vertices.get(i);
@@ -357,7 +358,6 @@ public class GraphBackedTypeStore implements ITypeStore {
 
         List<AtlasVertex> result = new ArrayList<>(infoList.size());
         List<String> typeNames = Lists.transform(infoList, new Function<TypeVertexInfo,String>() {
-
             @Override
             public String apply(TypeVertexInfo input) {
                 return input.getTypeName();

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/repository/typestore/StoreBackedTypeCache.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/repository/typestore/StoreBackedTypeCache.java b/repository/src/main/java/org/apache/atlas/repository/typestore/StoreBackedTypeCache.java
index 8573719..f472fa6 100644
--- a/repository/src/main/java/org/apache/atlas/repository/typestore/StoreBackedTypeCache.java
+++ b/repository/src/main/java/org/apache/atlas/repository/typestore/StoreBackedTypeCache.java
@@ -17,14 +17,9 @@
  */
 package org.apache.atlas.repository.typestore;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-
+import com.google.common.collect.ImmutableList;
 import org.apache.atlas.AtlasException;
+import org.apache.atlas.annotation.ConditionalOnAtlasProperty;
 import org.apache.atlas.typesystem.TypesDef;
 import org.apache.atlas.typesystem.types.AttributeDefinition;
 import org.apache.atlas.typesystem.types.ClassType;
@@ -38,10 +33,16 @@ import org.apache.atlas.typesystem.types.TypeSystem.TransientTypeSystem;
 import org.apache.atlas.typesystem.types.TypeUtils;
 import org.apache.atlas.typesystem.types.cache.DefaultTypeCache;
 import org.apache.atlas.typesystem.types.utils.TypesUtil;
+import org.springframework.stereotype.Component;
 
-import com.google.common.collect.ImmutableList;
-import com.google.inject.Inject;
-import com.google.inject.Singleton;
+import javax.inject.Inject;
+import javax.inject.Singleton;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 
 /**
@@ -52,7 +53,9 @@ import com.google.inject.Singleton;
  * are also loaded from the store if they are not already in the cache.
  */
 @Singleton
+@Component
 @Deprecated
+@ConditionalOnAtlasProperty(property = "atlas.TypeCache.impl")
 public class StoreBackedTypeCache extends DefaultTypeCache {
 
     private ITypeStore typeStore;

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java b/repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java
index fd85ef0..99d2107 100755
--- a/repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java
+++ b/repository/src/main/java/org/apache/atlas/services/DefaultMetadataService.java
@@ -20,8 +20,6 @@ package org.apache.atlas.services;
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
-import com.google.inject.Provider;
-import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.AtlasException;
@@ -35,6 +33,7 @@ import org.apache.atlas.listener.ChangedTypeDefs;
 import org.apache.atlas.listener.EntityChangeListener;
 import org.apache.atlas.listener.TypeDefChangeListener;
 import org.apache.atlas.listener.TypesChangeListener;
+import org.apache.atlas.model.legacy.EntityResult;
 import org.apache.atlas.repository.MetadataRepository;
 import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.audit.EntityAuditRepository;
@@ -54,16 +53,7 @@ import org.apache.atlas.typesystem.json.InstanceSerialization;
 import org.apache.atlas.typesystem.json.TypesSerialization;
 import org.apache.atlas.typesystem.persistence.Id;
 import org.apache.atlas.typesystem.persistence.ReferenceableInstance;
-import org.apache.atlas.typesystem.types.AttributeInfo;
-import org.apache.atlas.typesystem.types.ClassType;
-import org.apache.atlas.typesystem.types.DataTypes;
-import org.apache.atlas.typesystem.types.EnumTypeDefinition;
-import org.apache.atlas.typesystem.types.HierarchicalTypeDefinition;
-import org.apache.atlas.typesystem.types.IDataType;
-import org.apache.atlas.typesystem.types.Multiplicity;
-import org.apache.atlas.typesystem.types.StructTypeDefinition;
-import org.apache.atlas.typesystem.types.TraitType;
-import org.apache.atlas.typesystem.types.TypeSystem;
+import org.apache.atlas.typesystem.types.*;
 import org.apache.atlas.typesystem.types.cache.TypeCache;
 import org.apache.atlas.utils.ParamChecker;
 import org.apache.commons.configuration.Configuration;
@@ -71,6 +61,7 @@ import org.codehaus.jettison.json.JSONException;
 import org.codehaus.jettison.json.JSONObject;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
@@ -80,6 +71,7 @@ import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 
 
@@ -88,6 +80,7 @@ import java.util.Map;
  * for listening to changes to the repository.
  */
 @Singleton
+@Component
 @Deprecated
 public class DefaultMetadataService implements MetadataService, ActiveStateChangeHandler, TypeDefChangeListener {
     private enum OperationType {
@@ -106,24 +99,16 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
     private final Collection<TypesChangeListener> typeChangeListeners = new LinkedHashSet<>();
     private final Collection<EntityChangeListener> entityChangeListeners = new LinkedHashSet<>();
 
-    @Inject
     private EntityAuditRepository auditRepository;
 
     @Inject
-    DefaultMetadataService(final MetadataRepository repository, final ITypeStore typeStore,
-                           final Collection<Provider<TypesChangeListener>> typeListenerProviders,
-                           final Collection<Provider<EntityChangeListener>> entityListenerProviders, TypeCache typeCache)
-            throws AtlasException {
-        this(repository, typeStore, typeListenerProviders, entityListenerProviders,
-                TypeSystem.getInstance(), ApplicationProperties.get(), typeCache);
-    }
-    
-    //for testing only
     public DefaultMetadataService(final MetadataRepository repository, final ITypeStore typeStore,
-                           final Collection<Provider<TypesChangeListener>> typeListenerProviders,
-                           final Collection<Provider<EntityChangeListener>> entityListenerProviders,
-                           final TypeSystem typeSystem,
-                           final Configuration configuration, TypeCache typeCache) throws AtlasException {
+                                  final Set<TypesChangeListener> typesChangeListeners,
+                                  final Set<EntityChangeListener> entityChangeListeners,
+                                  final TypeSystem typeSystem,
+                                  final Configuration configuration,
+                                  TypeCache typeCache,
+                                  EntityAuditRepository auditRepository) throws AtlasException {
         this.typeStore = typeStore;
         this.typeSystem = typeSystem;
         /**
@@ -139,19 +124,17 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
 
         this.repository = repository;
 
-        for (Provider<TypesChangeListener> provider : typeListenerProviders) {
-            typeChangeListeners.add(provider.get());
-        }
+        this.typeChangeListeners.addAll(typesChangeListeners);
 
-        for (Provider<EntityChangeListener> provider : entityListenerProviders) {
-            entityChangeListeners.add(provider.get());
-        }
+        this.entityChangeListeners.addAll(entityChangeListeners);
 
         if (!HAConfiguration.isHAEnabled(configuration)) {
             restoreTypeSystem();
         }
 
         maxAuditResults = configuration.getShort(CONFIG_MAX_AUDIT_RESULTS, DEFAULT_MAX_AUDIT_RESULTS);
+
+        this.auditRepository = auditRepository;
     }
 
     private void restoreTypeSystem() throws AtlasException {
@@ -313,7 +296,7 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
     public ITypedReferenceableInstance[] deserializeClassInstances(String entityInstanceDefinition) throws AtlasException {
         return GraphHelper.deserializeClassInstances(typeSystem, entityInstanceDefinition);
     }
-    
+
     @Override
     public ITypedReferenceableInstance getTypedReferenceableInstance(Referenceable entityInstance) throws AtlasException {
         return GraphHelper.getTypedReferenceableInstance(typeSystem, entityInstance);
@@ -356,7 +339,7 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         return repository.getEntityDefinition(entityType, attribute, value);
     }
 
-        @Override
+    @Override
     public String getEntityDefinition(String entityType, String attribute, String value) throws AtlasException {
         final ITypedReferenceableInstance instance = getEntityDefinitionReference(entityType, attribute, value);
         return InstanceSerialization.toJson(instance, true);
@@ -422,7 +405,7 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         return result;
     }
 
-    private void onEntitiesAddedUpdated(AtlasClient.EntityResult entityResult) throws AtlasException {
+    private void onEntitiesAddedUpdated(EntityResult entityResult) throws AtlasException {
         onEntitiesAdded(entityResult.getCreatedEntities());
         onEntitiesUpdated(entityResult.getUpdateEntities());
         //Note: doesn't access deletedEntities from entityResult
@@ -635,8 +618,8 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
 
         // ensure trait is not already defined
         Preconditions
-            .checkArgument(!getTraitNames(guid).contains(traitName), "trait=%s is already defined for entity=%s",
-                    traitName, guid);
+                .checkArgument(!getTraitNames(guid).contains(traitName), "trait=%s is already defined for entity=%s",
+                        traitName, guid);
 
         repository.addTrait(guid, traitInstance);
 
@@ -644,7 +627,7 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
     }
 
     private ITypedStruct deserializeTraitInstance(String traitInstanceDefinition)
-    throws AtlasException {
+            throws AtlasException {
         return createTraitInstance(InstanceSerialization.fromJsonStruct(traitInstanceDefinition, true));
     }
 
@@ -762,14 +745,14 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
      * @see org.apache.atlas.services.MetadataService#deleteEntities(java.lang.String)
      */
     @Override
-    public AtlasClient.EntityResult deleteEntities(List<String> deleteCandidateGuids) throws AtlasException {
+    public EntityResult deleteEntities(List<String> deleteCandidateGuids) throws AtlasException {
         ParamChecker.notEmpty(deleteCandidateGuids, "delete candidate guids");
         return deleteGuids(deleteCandidateGuids);
     }
 
     @Override
-    public AtlasClient.EntityResult deleteEntityByUniqueAttribute(String typeName, String uniqueAttributeName,
-                                                                  String attrValue) throws AtlasException {
+    public EntityResult deleteEntityByUniqueAttribute(String typeName, String uniqueAttributeName,
+                                                      String attrValue) throws AtlasException {
         typeName            = ParamChecker.notEmpty(typeName, "delete candidate typeName");
         uniqueAttributeName = ParamChecker.notEmpty(uniqueAttributeName, "delete candidate unique attribute name");
         attrValue           = ParamChecker.notEmpty(attrValue, "delete candidate unique attribute value");
@@ -782,8 +765,8 @@ public class DefaultMetadataService implements MetadataService, ActiveStateChang
         return deleteGuids(deleteCandidateGuids);
     }
 
-    private AtlasClient.EntityResult deleteGuids(List<String> deleteCandidateGuids) throws AtlasException {
-        AtlasClient.EntityResult entityResult = repository.deleteEntities(deleteCandidateGuids);
+    private EntityResult deleteGuids(List<String> deleteCandidateGuids) throws AtlasException {
+        EntityResult entityResult = repository.deleteEntities(deleteCandidateGuids);
         onEntitiesAddedUpdated(entityResult);
         return entityResult;
     }

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/services/MetricsService.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/services/MetricsService.java b/repository/src/main/java/org/apache/atlas/services/MetricsService.java
index e3bedb8..0fa68b2 100644
--- a/repository/src/main/java/org/apache/atlas/services/MetricsService.java
+++ b/repository/src/main/java/org/apache/atlas/services/MetricsService.java
@@ -18,12 +18,11 @@
 package org.apache.atlas.services;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.inject.Singleton;
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasException;
+import org.apache.atlas.annotation.AtlasService;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.metrics.AtlasMetrics;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.util.AtlasGremlinQueryProvider;
 import org.apache.atlas.util.AtlasGremlinQueryProvider.AtlasGremlinQuery;
@@ -35,7 +34,7 @@ import javax.inject.Inject;
 import java.util.List;
 import java.util.Map;
 
-@Singleton
+@AtlasService
 public class MetricsService {
     private static final Logger LOG = LoggerFactory.getLogger(MetricsService.class);
 
@@ -75,8 +74,8 @@ public class MetricsService {
 
 
     @Inject
-    public MetricsService() throws AtlasException {
-        this(ApplicationProperties.get(), AtlasGraphProvider.getGraphInstance());
+    public MetricsService(AtlasGraph atlasGraph) throws AtlasException {
+        this(ApplicationProperties.get(), atlasGraph);
     }
 
     @VisibleForTesting

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/main/java/org/apache/atlas/util/AtlasRepositoryConfiguration.java
----------------------------------------------------------------------
diff --git a/repository/src/main/java/org/apache/atlas/util/AtlasRepositoryConfiguration.java b/repository/src/main/java/org/apache/atlas/util/AtlasRepositoryConfiguration.java
index a4a596a..09f366e 100644
--- a/repository/src/main/java/org/apache/atlas/util/AtlasRepositoryConfiguration.java
+++ b/repository/src/main/java/org/apache/atlas/util/AtlasRepositoryConfiguration.java
@@ -17,18 +17,13 @@
  */
 package org.apache.atlas.util;
 
-import java.util.ArrayList;
-import java.util.List;
-
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.repository.audit.EntityAuditRepository;
 import org.apache.atlas.repository.audit.HBaseBasedAuditRepository;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.DeleteHandler;
 import org.apache.atlas.repository.graph.SoftDeleteHandler;
 import org.apache.atlas.repository.graphdb.GraphDatabase;
-import org.apache.atlas.repository.graphdb.GremlinVersion;
 import org.apache.atlas.repository.store.graph.v1.DeleteHandlerV1;
 import org.apache.atlas.repository.store.graph.v1.SoftDeleteHandlerV1;
 import org.apache.atlas.typesystem.types.cache.DefaultTypeCache;
@@ -37,6 +32,9 @@ import org.apache.commons.configuration.Configuration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * Atlas configuration for repository project
  *

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/BaseRepositoryTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/BaseRepositoryTest.java b/repository/src/test/java/org/apache/atlas/BaseRepositoryTest.java
index a7cb2e5..7c6b60b 100644
--- a/repository/src/test/java/org/apache/atlas/BaseRepositoryTest.java
+++ b/repository/src/test/java/org/apache/atlas/BaseRepositoryTest.java
@@ -21,7 +21,6 @@ import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.repository.MetadataRepository;
-import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
 import org.apache.atlas.services.MetadataService;
 import org.apache.atlas.type.AtlasTypeRegistry;
@@ -45,7 +44,7 @@ import static org.apache.atlas.AtlasClient.PROCESS_ATTRIBUTE_OUTPUTS;
 /**
  *  Base Class to set up hive types and instances for tests
  */
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class BaseRepositoryTest {
 
     @Inject
@@ -70,7 +69,6 @@ public class BaseRepositoryTest {
 
     protected void tearDown() throws Exception {
         TypeSystem.getInstance().reset();
-        AtlasGraphProvider.cleanup();
     }
 
     private void setUpTypes() throws Exception {

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/RepositoryServiceLoadingTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/RepositoryServiceLoadingTest.java b/repository/src/test/java/org/apache/atlas/RepositoryServiceLoadingTest.java
index 01081da..506d11c 100755
--- a/repository/src/test/java/org/apache/atlas/RepositoryServiceLoadingTest.java
+++ b/repository/src/test/java/org/apache/atlas/RepositoryServiceLoadingTest.java
@@ -29,7 +29,7 @@ import org.testng.annotations.Test;
  * Uses TestNG's Guice annotation to load the necessary modules and inject the
  * objects from Guice
  */
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class RepositoryServiceLoadingTest {
 
     @Test

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/TestModules.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/TestModules.java b/repository/src/test/java/org/apache/atlas/TestModules.java
new file mode 100644
index 0000000..095af41
--- /dev/null
+++ b/repository/src/test/java/org/apache/atlas/TestModules.java
@@ -0,0 +1,231 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.atlas;
+
+import com.google.inject.Binder;
+import com.google.inject.Provider;
+import com.google.inject.Singleton;
+import com.google.inject.matcher.Matchers;
+import com.google.inject.multibindings.Multibinder;
+import org.apache.atlas.annotation.GraphTransaction;
+import org.apache.atlas.discovery.AtlasDiscoveryService;
+import org.apache.atlas.discovery.AtlasLineageService;
+import org.apache.atlas.discovery.DataSetLineageService;
+import org.apache.atlas.discovery.DiscoveryService;
+import org.apache.atlas.discovery.EntityDiscoveryService;
+import org.apache.atlas.discovery.EntityLineageService;
+import org.apache.atlas.discovery.LineageService;
+import org.apache.atlas.discovery.graph.GraphBackedDiscoveryService;
+import org.apache.atlas.graph.GraphSandboxUtil;
+import org.apache.atlas.listener.EntityChangeListener;
+import org.apache.atlas.listener.TypeDefChangeListener;
+import org.apache.atlas.listener.TypesChangeListener;
+import org.apache.atlas.repository.MetadataRepository;
+import org.apache.atlas.repository.audit.EntityAuditListener;
+import org.apache.atlas.repository.audit.EntityAuditRepository;
+import org.apache.atlas.repository.graph.DeleteHandler;
+import org.apache.atlas.repository.graph.GraphBackedMetadataRepository;
+import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
+import org.apache.atlas.repository.graph.HardDeleteHandler;
+import org.apache.atlas.repository.graph.SoftDeleteHandler;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
+import org.apache.atlas.repository.impexp.ExportService;
+import org.apache.atlas.repository.store.graph.AtlasEntityDefStore;
+import org.apache.atlas.repository.store.graph.AtlasEntityStore;
+import org.apache.atlas.repository.store.graph.v1.AtlasEntityChangeNotifier;
+import org.apache.atlas.repository.store.graph.v1.AtlasEntityDefStoreV1;
+import org.apache.atlas.repository.store.graph.v1.AtlasEntityStoreV1;
+import org.apache.atlas.repository.store.graph.v1.AtlasTypeDefGraphStoreV1;
+import org.apache.atlas.repository.store.graph.v1.DeleteHandlerV1;
+import org.apache.atlas.repository.store.graph.v1.EntityGraphMapper;
+import org.apache.atlas.repository.store.graph.v1.HardDeleteHandlerV1;
+import org.apache.atlas.repository.store.graph.v1.SoftDeleteHandlerV1;
+import org.apache.atlas.repository.typestore.GraphBackedTypeStore;
+import org.apache.atlas.repository.typestore.ITypeStore;
+import org.apache.atlas.repository.typestore.StoreBackedTypeCache;
+import org.apache.atlas.service.Service;
+import org.apache.atlas.services.DefaultMetadataService;
+import org.apache.atlas.services.MetadataService;
+import org.apache.atlas.store.AtlasTypeDefStore;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.typesystem.types.TypeSystem;
+import org.apache.atlas.typesystem.types.cache.TypeCache;
+import org.apache.atlas.util.AtlasRepositoryConfiguration;
+import org.apache.commons.configuration.Configuration;
+import org.mockito.Mockito;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class TestModules {
+
+    static class MockNotifier implements Provider<AtlasEntityChangeNotifier> {
+        @Override
+        public AtlasEntityChangeNotifier get() {
+            return Mockito.mock(AtlasEntityChangeNotifier.class);
+        }
+    }
+
+    // Test only DI modules
+    public static class TestOnlyModule extends com.google.inject.AbstractModule {
+
+        private static final Logger LOG = LoggerFactory.getLogger(TestOnlyModule.class);
+
+        static class TypeSystemProvider implements Provider<TypeSystem> {
+            @Override
+            public TypeSystem get() {
+                return TypeSystem.getInstance();
+            }
+        }
+
+        static class AtlasConfigurationProvider implements Provider<Configuration> {
+
+            @Override
+            public Configuration get() {
+                try {
+                    return ApplicationProperties.get();
+                } catch (AtlasException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }
+
+        static class AtlasGraphProvider implements Provider<AtlasGraph> {
+            @Override
+            public AtlasGraph get() {
+                return org.apache.atlas.repository.graph.AtlasGraphProvider.getGraphInstance();
+            }
+        }
+
+        @Override
+        protected void configure() {
+            GraphSandboxUtil.create();
+
+            bindAuditRepository(binder());
+
+            bindDeleteHandler(binder());
+
+            bind(AtlasGraph.class).toProvider(AtlasGraphProvider.class);
+
+            // allow for dynamic binding of the metadata repo & graph service
+            // bind the MetadataRepositoryService interface to an implementation
+            bind(MetadataRepository.class).to(GraphBackedMetadataRepository.class).asEagerSingleton();
+
+            bind(TypeSystem.class).toProvider(TypeSystemProvider.class).in(Singleton.class);
+            bind(Configuration.class).toProvider(AtlasConfigurationProvider.class).in(Singleton.class);
+
+            // bind the ITypeStore interface to an implementation
+            bind(ITypeStore.class).to(GraphBackedTypeStore.class).asEagerSingleton();
+            bind(AtlasTypeDefStore.class).to(AtlasTypeDefGraphStoreV1.class).asEagerSingleton();
+
+            //For testing
+            bind(AtlasEntityDefStore.class).to(AtlasEntityDefStoreV1.class).asEagerSingleton();
+            bind(AtlasTypeRegistry.class).asEagerSingleton();
+            bind(EntityGraphMapper.class).asEagerSingleton();
+            bind(ExportService.class).asEagerSingleton();
+
+            //GraphBackedSearchIndexer must be an eager singleton to force the search index creation to happen before
+            //we try to restore the type system (otherwise we'll end up running queries
+            //before we have any indices during the initial graph setup)
+            Multibinder<TypesChangeListener> typesChangeListenerBinder =
+                    Multibinder.newSetBinder(binder(), TypesChangeListener.class);
+            typesChangeListenerBinder.addBinding().to(GraphBackedSearchIndexer.class).asEagerSingleton();
+
+            // New typesdef/instance change listener should also be bound to the corresponding implementation
+            Multibinder<TypeDefChangeListener> typeDefChangeListenerMultibinder =
+                    Multibinder.newSetBinder(binder(), TypeDefChangeListener.class);
+            typeDefChangeListenerMultibinder.addBinding().to(DefaultMetadataService.class);
+            typeDefChangeListenerMultibinder.addBinding().to(GraphBackedSearchIndexer.class).asEagerSingleton();
+
+            bind(AtlasEntityStore.class).to(AtlasEntityStoreV1.class);
+
+            // bind the MetadataService interface to an implementation
+            bind(MetadataService.class).to(DefaultMetadataService.class).asEagerSingleton();
+
+            // bind the DiscoveryService interface to an implementation
+            bind(DiscoveryService.class).to(GraphBackedDiscoveryService.class).asEagerSingleton();
+            bind(AtlasDiscoveryService.class).to(EntityDiscoveryService.class).asEagerSingleton();
+
+            bind(LineageService.class).to(DataSetLineageService.class).asEagerSingleton();
+            bind(AtlasLineageService.class).to(EntityLineageService.class).asEagerSingleton();
+
+            bindTypeCache();
+
+            //Add EntityAuditListener as EntityChangeListener
+            Multibinder<EntityChangeListener> entityChangeListenerBinder =
+                    Multibinder.newSetBinder(binder(), EntityChangeListener.class);
+            entityChangeListenerBinder.addBinding().to(EntityAuditListener.class);
+
+            final GraphTransactionInterceptor graphTransactionInterceptor = new GraphTransactionInterceptor(new AtlasGraphProvider().get());
+            requestInjection(graphTransactionInterceptor);
+            bindInterceptor(Matchers.any(), Matchers.annotatedWith(GraphTransaction.class), graphTransactionInterceptor);
+        }
+
+        protected void bindTypeCache() {
+            bind(TypeCache.class).to(AtlasRepositoryConfiguration.getTypeCache()).asEagerSingleton();
+        }
+
+        protected void bindDeleteHandler(Binder binder) {
+            binder.bind(DeleteHandler.class).to(AtlasRepositoryConfiguration.getDeleteHandlerImpl()).asEagerSingleton();
+            binder.bind(DeleteHandlerV1.class).to(AtlasRepositoryConfiguration.getDeleteHandlerV1Impl()).asEagerSingleton();
+        }
+
+        protected void bindAuditRepository(Binder binder) {
+
+            Class<? extends EntityAuditRepository> auditRepoImpl = AtlasRepositoryConfiguration.getAuditRepositoryImpl();
+
+            //Map EntityAuditRepository interface to configured implementation
+            binder.bind(EntityAuditRepository.class).to(auditRepoImpl).asEagerSingleton();
+
+            if(Service.class.isAssignableFrom(auditRepoImpl)) {
+                Class<? extends Service> auditRepoService = (Class<? extends Service>)auditRepoImpl;
+                //if it's a service, make sure that it gets properly closed at shutdown
+                Multibinder<Service> serviceBinder = Multibinder.newSetBinder(binder, Service.class);
+                serviceBinder.addBinding().to(auditRepoService);
+            }
+        }
+    }
+
+    public static class SoftDeleteModule extends TestOnlyModule {
+        @Override
+        protected void bindDeleteHandler(Binder binder) {
+            bind(DeleteHandler.class).to(SoftDeleteHandler.class).asEagerSingleton();
+            bind(DeleteHandlerV1.class).to(SoftDeleteHandlerV1.class).asEagerSingleton();
+            bind(AtlasEntityChangeNotifier.class).toProvider(MockNotifier.class);
+        }
+    }
+
+    public static class HardDeleteModule extends TestOnlyModule {
+        @Override
+        protected void bindDeleteHandler(Binder binder) {
+            bind(DeleteHandler.class).to(HardDeleteHandler.class).asEagerSingleton();
+            bind(DeleteHandlerV1.class).to(HardDeleteHandlerV1.class).asEagerSingleton();
+            bind(AtlasEntityChangeNotifier.class).toProvider(MockNotifier.class);
+        }
+    }
+
+    /**
+     * Guice module which sets TypeCache implementation class configuration property to {@link StoreBackedTypeCache}.
+     *
+     */
+    public static class StoreBackedTypeCacheTestModule extends TestOnlyModule {
+        @Override
+        protected void bindTypeCache() {
+            bind(TypeCache.class).to(StoreBackedTypeCache.class).asEagerSingleton();
+        }
+    }
+}

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/TestOnlyModule.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/TestOnlyModule.java b/repository/src/test/java/org/apache/atlas/TestOnlyModule.java
deleted file mode 100644
index c930953..0000000
--- a/repository/src/test/java/org/apache/atlas/TestOnlyModule.java
+++ /dev/null
@@ -1,30 +0,0 @@
-/**
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- * <p>
- * http://www.apache.org/licenses/LICENSE-2.0
- * <p>
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.atlas;
-
-import org.apache.atlas.graph.GraphSandboxUtil;
-
-public class TestOnlyModule extends RepositoryMetadataModule {
-    @Override
-    protected void configure() {
-        GraphSandboxUtil.create();
-
-        // Configure extra stuff for test DI here
-        super.configure();
-    }
-}

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/TestUtils.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/TestUtils.java b/repository/src/test/java/org/apache/atlas/TestUtils.java
index cf39d8d..1e439f7 100755
--- a/repository/src/test/java/org/apache/atlas/TestUtils.java
+++ b/repository/src/test/java/org/apache/atlas/TestUtils.java
@@ -20,10 +20,11 @@ package org.apache.atlas;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.google.inject.Provider;
+import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.listener.EntityChangeListener;
 import org.apache.atlas.listener.TypesChangeListener;
 import org.apache.atlas.repository.MetadataRepository;
+import org.apache.atlas.repository.audit.InMemoryEntityAuditRepository;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphBackedMetadataRepository;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
@@ -64,7 +65,6 @@ import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -600,20 +600,18 @@ public final class TestUtils {
             typeCache = new DefaultTypeCache();
         }
         final GraphBackedSearchIndexer indexer = new GraphBackedSearchIndexer(new AtlasTypeRegistry());
-        Provider<TypesChangeListener> indexerProvider = new Provider<TypesChangeListener>() {
-
-            @Override
-            public TypesChangeListener get() {
-                return indexer;
-            }
-        };
 
         Configuration config = ApplicationProperties.get();
-        ITypeStore typeStore = new GraphBackedTypeStore();
+        ITypeStore typeStore = new GraphBackedTypeStore(AtlasGraphProvider.getGraphInstance());
         DefaultMetadataService defaultMetadataService = new DefaultMetadataService(repo,
                 typeStore,
-                Collections.singletonList(indexerProvider),
-                new ArrayList<Provider<EntityChangeListener>>(), TypeSystem.getInstance(), config, typeCache);
+                new HashSet<TypesChangeListener>() {{ add(indexer); }},
+                new HashSet<EntityChangeListener>(),
+                TypeSystem.getInstance(),
+                config,
+                typeCache,
+                // Fixme: Can we work with Noop
+                new InMemoryEntityAuditRepository());
 
         //commit the created types
         getGraph().commit();

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/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 fc38edf..c700a61 100644
--- a/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
+++ b/repository/src/test/java/org/apache/atlas/discovery/DataSetLineageServiceTest.java
@@ -19,10 +19,10 @@
 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.TestOnlyModule;
+import org.apache.atlas.TestModules;
+import org.apache.atlas.model.legacy.EntityResult;
 import org.apache.atlas.query.QueryParams;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.Referenceable;
@@ -55,7 +55,7 @@ import static org.testng.Assert.fail;
 /**
  * Unit tests for Hive LineageService.
  */
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class DataSetLineageServiceTest extends BaseRepositoryTest {
 
     @Inject
@@ -400,7 +400,7 @@ public class DataSetLineageServiceTest extends BaseRepositoryTest {
 
         //Delete the entity. Lineage for entity returns the same results as before.
         //Lineage for table name throws EntityNotFoundException
-        AtlasClient.EntityResult deleteResult = repository.deleteEntities(Arrays.asList(tableId));
+        EntityResult deleteResult = repository.deleteEntities(Arrays.asList(tableId));
         assertTrue(deleteResult.getDeletedEntities().contains(tableId));
 
         results = new JSONObject(lineageService.getSchemaForEntity(tableId));

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java b/repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java
index 3ea8f34..834abe1 100755
--- a/repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java
+++ b/repository/src/test/java/org/apache/atlas/discovery/GraphBackedDiscoveryServiceTest.java
@@ -22,7 +22,7 @@ import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.BaseRepositoryTest;
 import org.apache.atlas.RequestContext;
-import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtils;
 import org.apache.atlas.discovery.graph.GraphBackedDiscoveryService;
 import org.apache.atlas.query.QueryParams;
@@ -62,7 +62,7 @@ import static org.apache.atlas.typesystem.types.utils.TypesUtil.createRequiredAt
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class GraphBackedDiscoveryServiceTest extends BaseRepositoryTest {
 
     @Inject

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/lineage/EntityLineageServiceTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/lineage/EntityLineageServiceTest.java b/repository/src/test/java/org/apache/atlas/lineage/EntityLineageServiceTest.java
index 0a4a445..8a88f32 100644
--- a/repository/src/test/java/org/apache/atlas/lineage/EntityLineageServiceTest.java
+++ b/repository/src/test/java/org/apache/atlas/lineage/EntityLineageServiceTest.java
@@ -19,15 +19,15 @@
 package org.apache.atlas.lineage;
 
 import com.google.common.collect.ImmutableList;
-import org.apache.atlas.AtlasClient;
 import org.apache.atlas.AtlasErrorCode;
 import org.apache.atlas.BaseRepositoryTest;
-import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtils;
 import org.apache.atlas.discovery.EntityLineageService;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasEntity.Status;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
+import org.apache.atlas.model.legacy.EntityResult;
 import org.apache.atlas.model.lineage.AtlasLineageInfo;
 import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageDirection;
 import org.apache.atlas.model.lineage.AtlasLineageInfo.LineageRelation;
@@ -56,7 +56,7 @@ import static org.testng.Assert.fail;
 /**
  * Unit tests for the new v2 Instance LineageService.
  */
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class EntityLineageServiceTest extends BaseRepositoryTest {
 
     @Inject
@@ -301,7 +301,7 @@ public class EntityLineageServiceTest extends BaseRepositoryTest {
 
         //Delete the table entity. Lineage for entity returns the same results as before.
         //Lineage for table name throws EntityNotFoundException
-        AtlasClient.EntityResult deleteResult = repository.deleteEntities(Arrays.asList(entityGuid));
+        EntityResult deleteResult = repository.deleteEntities(Arrays.asList(entityGuid));
         assertTrue(deleteResult.getDeletedEntities().contains(entityGuid));
 
         inputLineage = getInputLineageInfo(entityGuid, 5);

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/AbstractGremlinQueryOptimizerTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/AbstractGremlinQueryOptimizerTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/AbstractGremlinQueryOptimizerTest.java
index 59d2a3f..88de2c4 100644
--- a/repository/src/test/java/org/apache/atlas/repository/graph/AbstractGremlinQueryOptimizerTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/AbstractGremlinQueryOptimizerTest.java
@@ -63,12 +63,12 @@ public abstract class AbstractGremlinQueryOptimizerTest implements IAtlasGraphPr
     private final GraphPersistenceStrategies STRATEGY = mock(GraphPersistenceStrategies.class);
 
     @BeforeClass
-    public void setUp() {
+    public void setUp() throws RepositoryException {
         GremlinQueryOptimizer.reset();
         GremlinQueryOptimizer.setExpressionFactory(getFactory());
         when(STRATEGY.typeAttributeName()).thenReturn(Constants.ENTITY_TYPE_PROPERTY_KEY);
         when(STRATEGY.superTypeAttributeName()).thenReturn(Constants.SUPER_TYPES_PROPERTY_KEY);
-        repo = new GraphBackedMetadataRepository(this, new HardDeleteHandler(TypeSystem.getInstance()));
+        repo = new GraphBackedMetadataRepository(new HardDeleteHandler(TypeSystem.getInstance()), this.get());
     }
 
     private FieldInfo getTestFieldInfo() throws AtlasException {

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java b/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java
index 4fd416c..0e22080 100644
--- a/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryDeleteTestBase.java
@@ -21,12 +21,12 @@ package org.apache.atlas.repository.graph;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.AtlasClient;
-import org.apache.atlas.AtlasClient.EntityResult;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.CreateUpdateEntitiesResult;
 import org.apache.atlas.RequestContext;
-import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtils;
+import org.apache.atlas.model.legacy.EntityResult;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.MetadataRepository;
 import org.apache.atlas.repository.RepositoryException;
@@ -53,6 +53,7 @@ import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
+import javax.inject.Inject;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -72,7 +73,7 @@ import static org.testng.Assert.fail;
  * Guice loads the dependencies and injects the necessary objects
  *
  */
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
 
     protected MetadataRepository repositoryService;
@@ -83,6 +84,9 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
 
     private ClassType compositeMapValueType;
 
+    @Inject
+    AtlasGraph atlasGraph;
+
     @BeforeClass
     public void setUp() throws Exception {
 
@@ -90,7 +94,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         typeSystem.reset();
 
         new GraphBackedSearchIndexer(new AtlasTypeRegistry());
-        final GraphBackedMetadataRepository delegate = new GraphBackedMetadataRepository(getDeleteHandler(typeSystem));
+        final GraphBackedMetadataRepository delegate = new GraphBackedMetadataRepository(getDeleteHandler(typeSystem), atlasGraph);
 
         repositoryService = TestUtils.addTransactionWrapper(delegate);
 
@@ -126,7 +130,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
     @AfterClass
     public void tearDown() throws Exception {
         TypeSystem.getInstance().reset();
-        AtlasGraphProvider.cleanup();
+//        AtlasGraphProvider.cleanup();
     }
 
     @Test
@@ -194,7 +198,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         table.set(COLUMNS_ATTR_NAME, Arrays.asList(new Id(colId, 0, COLUMN_TYPE)));
         String tableId = createInstance(table);
 
-        AtlasClient.EntityResult entityResult = deleteEntities(colId);
+        EntityResult entityResult = deleteEntities(colId);
         assertEquals(entityResult.getDeletedEntities().size(), 1);
         assertEquals(entityResult.getDeletedEntities().get(0), colId);
         assertEquals(entityResult.getUpdateEntities().size(), 1);
@@ -227,7 +231,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
 
     protected abstract void assertEntityDeleted(String id) throws Exception;
 
-    private AtlasClient.EntityResult deleteEntities(String... id) throws Exception {
+    private EntityResult deleteEntities(String... id) throws Exception {
         RequestContext.createContext();
         return repositoryService.deleteEntities(Arrays.asList(id));
     }
@@ -261,7 +265,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         String colId = columns.get(0).getId()._getId();
         String tableId = tableInstance.getId()._getId();
 
-        AtlasClient.EntityResult entityResult = deleteEntities(colId);
+        EntityResult entityResult = deleteEntities(colId);
         assertEquals(entityResult.getDeletedEntities().size(), 1);
         assertEquals(entityResult.getDeletedEntities().get(0), colId);
         assertEquals(entityResult.getUpdateEntities().size(), 1);
@@ -389,7 +393,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         return entityDefinition;
     }
 
-    private AtlasClient.EntityResult updatePartial(ITypedReferenceableInstance entity) throws RepositoryException {
+    private EntityResult updatePartial(ITypedReferenceableInstance entity) throws RepositoryException {
         RequestContext.createContext();
         return repositoryService.updatePartial(entity).getEntityResult();
     }
@@ -419,7 +423,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         ClassType personType = typeSystem.getDataType(ClassType.class, "Person");
         ITypedReferenceableInstance maxEntity = personType.createInstance(max.getId());
         maxEntity.set("mentor", johnGuid);
-        AtlasClient.EntityResult entityResult = updatePartial(maxEntity);
+        EntityResult entityResult = updatePartial(maxEntity);
         assertEquals(entityResult.getUpdateEntities().size(), 1);
         assertTrue(entityResult.getUpdateEntities().contains(maxGuid));
 
@@ -512,7 +516,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         Assert.assertTrue(subordinateIds.contains(maxGuid));
 
 
-        AtlasClient.EntityResult entityResult = deleteEntities(maxGuid);
+        EntityResult entityResult = deleteEntities(maxGuid);
         ITypedReferenceableInstance john = repositoryService.getEntityDefinition("Person", "name", "John");
 
         assertEquals(entityResult.getDeletedEntities().size(), 1);
@@ -560,7 +564,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         String columnGuid = columns.get(0).getId()._getId();
 
         // Delete the column.
-        AtlasClient.EntityResult entityResult = deleteEntities(columnGuid);
+        EntityResult entityResult = deleteEntities(columnGuid);
         assertEquals(entityResult.getDeletedEntities().size(), 1);
         Assert.assertTrue(entityResult.getDeletedEntities().contains(columnGuid));
         assertEquals(entityResult.getUpdateEntities().size(), 1);
@@ -672,7 +676,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         Assert.assertEquals(refList.get(0).getId()._getId(), traitTargetGuid);
 
         // Delete the entities that are targets of the struct and trait instances.
-        AtlasClient.EntityResult entityResult = deleteEntities(structTargetGuid, traitTargetGuid);
+        EntityResult entityResult = deleteEntities(structTargetGuid, traitTargetGuid);
         Assert.assertEquals(entityResult.getDeletedEntities().size(), 2);
         Assert.assertTrue(entityResult.getDeletedEntities().containsAll(Arrays.asList(structTargetGuid, traitTargetGuid)));
         assertEntityDeleted(structTargetGuid);
@@ -1037,7 +1041,7 @@ public abstract class GraphBackedMetadataRepositoryDeleteTestBase {
         String table2Id = createInstance(table2);
 
         // Delete the tables and column
-        AtlasClient.EntityResult entityResult = deleteEntities(table1Id, colId, table2Id);
+        EntityResult entityResult = deleteEntities(table1Id, colId, table2Id);
         Assert.assertEquals(entityResult.getDeletedEntities().size(), 3);
         Assert.assertTrue(entityResult.getDeletedEntities().containsAll(Arrays.asList(colId, table1Id, table2Id)));
         assertEntityDeleted(table1Id);

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryTest.java
index 9f44c31..8120aaa 100755
--- a/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedMetadataRepositoryTest.java
@@ -22,10 +22,10 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.CreateUpdateEntitiesResult;
-import org.apache.atlas.GraphTransaction;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.RequestContext;
-import org.apache.atlas.TestOnlyModule;
 import org.apache.atlas.TestUtils;
+import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.discovery.graph.GraphBackedDiscoveryService;
 import org.apache.atlas.query.QueryParams;
 import org.apache.atlas.repository.Constants;
@@ -92,7 +92,7 @@ import static org.testng.Assert.assertTrue;
  * Guice loads the dependencies and injects the necessary objects
  *
  */
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class GraphBackedMetadataRepositoryTest {
 
     @Inject
@@ -126,7 +126,7 @@ public class GraphBackedMetadataRepositoryTest {
     @AfterClass
     public void tearDown() throws Exception {
         TypeSystem.getInstance().reset();
-        AtlasGraphProvider.cleanup();
+//        AtlasGraphProvider.cleanup();
     }
 
     @Test

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexerTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexerTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexerTest.java
index a254d0a..ef5d469 100644
--- a/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexerTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/GraphBackedSearchIndexerTest.java
@@ -20,7 +20,7 @@ package org.apache.atlas.repository.graph;
 
 import com.google.inject.Inject;
 import org.apache.atlas.AtlasException;
-import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtils;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
@@ -41,12 +41,10 @@ import org.testng.annotations.Test;
 import java.util.Arrays;
 import java.util.Set;
 
-import static junit.framework.Assert.assertTrue;
 import static org.apache.atlas.typesystem.types.utils.TypesUtil.createClassTypeDef;
-import static org.testng.Assert.assertFalse;
-import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.*;
 
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class GraphBackedSearchIndexerTest {
     @Inject
     private GraphBackedSearchIndexer graphBackedSearchIndexer;

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/GraphHelperTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/GraphHelperTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/GraphHelperTest.java
index 4061543..e64c2c8 100644
--- a/repository/src/test/java/org/apache/atlas/repository/graph/GraphHelperTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/GraphHelperTest.java
@@ -19,7 +19,7 @@
 package org.apache.atlas.repository.graph;
 
 import org.apache.atlas.AtlasException;
-import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtils;
 import org.apache.atlas.repository.graph.GraphHelper.VertexInfo;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
@@ -56,7 +56,7 @@ import java.util.Set;
 
 import static org.testng.Assert.*;
 
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class GraphHelperTest {
 
 
@@ -104,7 +104,7 @@ public class GraphHelperTest {
 
     @AfterClass
     public void tearDown() {
-        AtlasGraphProvider.cleanup();
+//        AtlasGraphProvider.cleanup();
     }
 
     @Test

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/GraphRepoMapperScaleTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/GraphRepoMapperScaleTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/GraphRepoMapperScaleTest.java
index 21926dc..1f077ba 100755
--- a/repository/src/test/java/org/apache/atlas/repository/graph/GraphRepoMapperScaleTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/GraphRepoMapperScaleTest.java
@@ -20,9 +20,9 @@ package org.apache.atlas.repository.graph;
 
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.CreateUpdateEntitiesResult;
-import org.apache.atlas.GraphTransaction;
-import org.apache.atlas.TestOnlyModule;
+import org.apache.atlas.TestModules;
 import org.apache.atlas.TestUtils;
+import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
@@ -54,7 +54,7 @@ import java.util.Date;
 import java.util.Iterator;
 
 @Test
-@Guice(modules = TestOnlyModule.class)
+@Guice(modules = TestModules.TestOnlyModule.class)
 public class GraphRepoMapperScaleTest {
 
     private static final String DATABASE_NAME = "foo";
@@ -89,7 +89,7 @@ public class GraphRepoMapperScaleTest {
     @AfterClass
     public void tearDown() throws Exception {
         TypeSystem.getInstance().reset();
-        AtlasGraphProvider.cleanup();
+//        AtlasGraphProvider.cleanup();
     }
 
     @Test

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateHardDeleteTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateHardDeleteTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateHardDeleteTest.java
index 9778c63..a3207fb 100644
--- a/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateHardDeleteTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateHardDeleteTest.java
@@ -17,9 +17,10 @@
  */
 package org.apache.atlas.repository.graph;
 
+import org.apache.atlas.TestModules;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
-import org.apache.atlas.typesystem.types.TypeSystem;
 import org.testng.Assert;
+import org.testng.annotations.Guice;
 
 import java.util.List;
 
@@ -28,14 +29,8 @@ import java.util.List;
  * Run tests in {@link ReverseReferenceUpdateTestBase} with hard delete enabled.
  *
  */
+@Guice(modules = TestModules.HardDeleteModule.class)
 public class ReverseReferenceUpdateHardDeleteTest extends ReverseReferenceUpdateTestBase {
-
-    @Override
-    DeleteHandler getDeleteHandler(TypeSystem typeSystem) {
-
-        return new HardDeleteHandler(typeSystem);
-    }
-
     @Override
     void assertTestOneToOneReference(Object refValue, ITypedReferenceableInstance expectedValue, ITypedReferenceableInstance referencingInstance) throws Exception {
         // Verify reference was disconnected
@@ -48,5 +43,4 @@ public class ReverseReferenceUpdateHardDeleteTest extends ReverseReferenceUpdate
         List<ITypedReferenceableInstance> refValues = (List<ITypedReferenceableInstance>) object;
         Assert.assertEquals(refValues.size(), 1);
     }
-
 }

http://git-wip-us.apache.org/repos/asf/incubator-atlas/blob/35e5828f/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateSoftDeleteTest.java
----------------------------------------------------------------------
diff --git a/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateSoftDeleteTest.java b/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateSoftDeleteTest.java
index 100ccfb..c5eda37 100644
--- a/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateSoftDeleteTest.java
+++ b/repository/src/test/java/org/apache/atlas/repository/graph/ReverseReferenceUpdateSoftDeleteTest.java
@@ -17,13 +17,14 @@
  */
 package org.apache.atlas.repository.graph;
 
+import org.apache.atlas.TestModules;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.typesystem.ITypedReferenceableInstance;
 import org.apache.atlas.typesystem.persistence.Id;
-import org.apache.atlas.typesystem.types.TypeSystem;
 import org.testng.Assert;
+import org.testng.annotations.Guice;
 
 import java.util.Iterator;
 import java.util.List;
@@ -33,14 +34,8 @@ import java.util.List;
  * Run tests in {@link ReverseReferenceUpdateTestBase} with soft delete enabled.
  *
  */
+@Guice(modules = TestModules.SoftDeleteModule.class)
 public class ReverseReferenceUpdateSoftDeleteTest extends ReverseReferenceUpdateTestBase {
-
-    @Override
-    DeleteHandler getDeleteHandler(TypeSystem typeSystem) {
-
-        return new SoftDeleteHandler(typeSystem);
-    }
-
     @Override
     void assertTestOneToOneReference(Object actual, ITypedReferenceableInstance expectedValue, ITypedReferenceableInstance referencingInstance) throws Exception {
         // Verify reference was not disconnected if soft deletes are enabled.