Retaining context in APEX Engine based on policies received in PdpUpdate

Change-Id: I73fad5bf76ed6b4979f5ab76013f204ea82da30b
Issue-ID: POLICY-2215
Signed-off-by: a.sreekumar <ajith.sreekumar@est.tech>

2 files changed, 7 insertions, 10 deletions

diff --git a/context/context-management/src/main/java/org/onap/policy/apex/context/Distributor.java b/context/context-management/src/main/java/org/onap/policy/apex/context/Distributor.java
index a17313882..430832382 100644
--- a/context/context-management/src/main/java/org/onap/policy/apex/context/Distributor.java
+++ b/context/context-management/src/main/java/org/onap/policy/apex/context/Distributor.java
@@ -1,6 +1,7 @@
 /*-
  * ============LICENSE_START=======================================================
  *  Copyright (C) 2016-2018 Ericsson. All rights reserved.
+ *  Modifications Copyright (C) 2019 Nordix Foundation.
  * ================================================================================
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -75,10 +76,10 @@ public interface Distributor {
     /**
      * Remove a context album from a distributor.
      *
-     * @param contextAlbum The album to remove
+     * @param axContextAlbumKey The key of album to remove
      * @throws ContextException if the album cannot be removed
      */
-    void removeContextAlbum(AxContextAlbum contextAlbum) throws ContextException;
+    void removeContextAlbum(AxArtifactKey axContextAlbumKey) throws ContextException;
 
     /**
      * Flush all context albums owned by the distributor to the distribution mechanism.
diff --git a/context/context-management/src/main/java/org/onap/policy/apex/context/impl/distribution/AbstractDistributor.java b/context/context-management/src/main/java/org/onap/policy/apex/context/impl/distribution/AbstractDistributor.java
index 49b0a669a..1d73f58f1 100644
--- a/context/context-management/src/main/java/org/onap/policy/apex/context/impl/distribution/AbstractDistributor.java
+++ b/context/context-management/src/main/java/org/onap/policy/apex/context/impl/distribution/AbstractDistributor.java
@@ -1,6 +1,7 @@
 /*-
  * ============LICENSE_START=======================================================
  *  Copyright (C) 2016-2018 Ericsson. All rights reserved.
+ *  Modifications Copyright (C) 2019 Nordix Foundation.
  * ================================================================================
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,7 +25,6 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
-
 import org.onap.policy.apex.context.ContextAlbum;
 import org.onap.policy.apex.context.ContextException;
 import org.onap.policy.apex.context.Distributor;
@@ -210,16 +210,12 @@ public abstract class AbstractDistributor implements Distributor {
      * {@inheritDoc}.
      */
     @Override
-    public void removeContextAlbum(final AxContextAlbum contextAlbum) throws ContextException {
+    public void removeContextAlbum(final AxArtifactKey axContextAlbumKey) throws ContextException {
         synchronized (albumMaps) {
-            // Check if the map already exists, if not return
-            if (!albumMaps.containsKey(contextAlbum.getKey())) {
-                LOGGER.warn("map remove failed, supplied map is null");
+            // Remove the map from the distributor
+            if (null == albumMaps.remove(axContextAlbumKey)) {
                 throw new ContextException("map update failed, supplied map is null");
             }
-
-            // Remove the map from the distributor
-            albumMaps.remove(contextAlbum.getKey());
         }
     }