4 files changed, 75 insertions, 13 deletions

diff --git a/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/DistributedLockingFeature.java b/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/DistributedLockingFeature.java
index cc7a7a12..3d19c873 100644
--- a/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/DistributedLockingFeature.java
+++ b/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/DistributedLockingFeature.java
@@ -65,9 +65,12 @@ public class DistributedLockingFeature implements PolicyEngineFeatureAPI, Policy
 	 */
 	private static final UUID uuid = UUID.randomUUID();
 	
+
 	/**
-	 * Config directory
+	 * Reference to Heartbeat
 	 */
+	private static Heartbeat heartbeat = null;
+	
 	@Override
 	public int getSequenceNumber() {
         return 1000;
@@ -116,7 +119,7 @@ public class DistributedLockingFeature implements PolicyEngineFeatureAPI, Policy
 		long heartbeatInterval = this.lockProps.getHeartBeatIntervalProperty();
 		
 		cleanLockTable();
-		Heartbeat heartbeat = new Heartbeat(this.uuid, lockProps);
+		heartbeat = new Heartbeat(this.uuid, lockProps);
 		
 		this.scheduledExecutorService = Executors.newScheduledThreadPool(1);
 		this.scheduledExecutorService.scheduleAtFixedRate(heartbeat, heartbeatInterval, heartbeatInterval, TimeUnit.MILLISECONDS);
@@ -154,5 +157,9 @@ public class DistributedLockingFeature implements PolicyEngineFeatureAPI, Policy
 		}
 		
 	}
+
+	public static Heartbeat getHeartbeat() {
+		return heartbeat;
+	}
 	
 }
diff --git a/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/Heartbeat.java b/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/Heartbeat.java
index c753dba9..edd0782e 100644
--- a/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/Heartbeat.java
+++ b/feature-distributed-locking/src/main/java/org/onap/policy/distributed/locking/Heartbeat.java
@@ -24,10 +24,9 @@ import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
-import java.util.Properties;
 import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
 
-import org.onap.policy.drools.utils.NetworkUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -52,14 +51,34 @@ public class Heartbeat implements Runnable{
 	 */
 	private UUID uuid;
 	
+	/**
+	 * Countdown latch for testing 
+	 */
+	private volatile CountDownLatch latch;
+	
+	/**
+	 * 
+	 * @param uuid
+	 * @param lockProps
+	 */
 	public Heartbeat(UUID uuid, DistributedLockingProperties lockProps) {
 		this.lockProps = lockProps;
 		this.uuid = uuid;
+		this.latch = new CountDownLatch(1);
+	}
+	/**
+	 * 
+	 * @param latch
+	 * Used for testing purposes
+	 */
+	protected void giveLatch(CountDownLatch latch) {
+		this.latch = latch;
 	}
 	
 	@Override
 	public void run() {
 
+		
 		long expirationAge = lockProps.getAgingProperty();
 
 		try (Connection conn = DriverManager.getConnection(lockProps.getDbUrl(), lockProps.getDbUser(),
@@ -70,9 +89,14 @@ public class Heartbeat implements Runnable{
 			statement.setLong(1, System.currentTimeMillis() + expirationAge);
 			statement.setString(2, this.uuid.toString());
 			statement.executeUpdate();
+			
+			latch.countDown();
+			
 		} catch (SQLException e) {
 			logger.error("error in Heartbeat.run()", e);
 		}
 
 	}
+	
+	
 }
diff --git a/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/test/DistributedLockingFeatureExceptionTest.java b/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/DistributedLockingFeatureExceptionTest.java
index ea53e522..5bc3ae5e 100644
--- a/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/test/DistributedLockingFeatureExceptionTest.java
+++ b/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/DistributedLockingFeatureExceptionTest.java
@@ -18,7 +18,7 @@
  * ============LICENSE_END=========================================================
  */
 
-package org.onap.policy.distributed.locking.test;
+package org.onap.policy.distributed.locking;
 
 import static org.junit.Assert.assertEquals;
 
diff --git a/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/test/TargetLockTest.java b/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/TargetLockTest.java
index e624afb9..2dc2ceb8 100644
--- a/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/test/TargetLockTest.java
+++ b/feature-distributed-locking/src/test/java/org/onap/policy/distributed/locking/TargetLockTest.java
@@ -18,7 +18,7 @@
  * ============LICENSE_END=========================================================
  */
 
-package org.onap.policy.distributed.locking.test;
+package org.onap.policy.distributed.locking;
 
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -70,7 +70,6 @@ public class TargetLockTest {
 	
 	@Before
 	public void wipeDb() {
-		
 		try (PreparedStatement lockDelete = conn.prepareStatement("DELETE FROM pooling.locks");){
 			lockDelete.executeUpdate();
 		} catch (SQLException e) {
@@ -101,16 +100,31 @@ public class TargetLockTest {
 
 	@Test
 	public void testExpiredLocks() throws InterruptedException, ExecutionException {
-		CountDownLatch latch = new CountDownLatch(1);
-		
+
+		CountDownLatch rowExpireLatch = new CountDownLatch(1);
+		CountDownLatch heartbeatLatch = new CountDownLatch(1);
+			
+			//grab lock
 		distLockFeat.beforeLock("resource1", "owner1", null);
 		
+			//Wait for lock to expire
 		try {
-			latch.await(1000, TimeUnit.MILLISECONDS);
+			rowExpireLatch.await(150, TimeUnit.MILLISECONDS);
 		} catch (InterruptedException e) {
 			logger.error("Error in testExpiredLocks", e);
 		}
 		
+		// Grab reference to heartbeat object
+		Heartbeat heartbeat = distLockFeat.getHeartbeat();
+
+		// Pass heartbeat object countdown latch
+		try {
+			heartbeat.giveLatch(heartbeatLatch);
+			heartbeatLatch.await(5, TimeUnit.SECONDS);
+		} catch (InterruptedException e) {
+			logger.error("Error in testExpiredLocks", e);
+		}
+	
 			//Heartbeat should keep it active
 		assertFalse(distLockFeat.beforeLock("resource1", "owner1", null).get());
 	}
@@ -169,16 +183,33 @@ public class TargetLockTest {
 	
 	@Test
 	public void testHeartbeat() {
-		CountDownLatch latch = new CountDownLatch(1);
 		
+		CountDownLatch rowExpireLatch = new CountDownLatch(1);
+		CountDownLatch heartbeatLatch = new CountDownLatch(1);
+		
+			//grab lock
 		distLockFeat.beforeLock("resource1", "owner1", null);
+		
+			//Wait for lock to expire
 		try {
-			latch.await(1000, TimeUnit.MILLISECONDS);
+			rowExpireLatch.await(150, TimeUnit.MILLISECONDS);
 		} catch (InterruptedException e) {
 			logger.error("Error in testExpiredLocks", e);
 		}
 		
-			// This test always returns true.
+			//Grab reference to heartbeat object
+		Heartbeat heartbeat = distLockFeat.getHeartbeat();
+		
+			//Pass heartbeat object countdown latch
+		try {
+			heartbeat.giveLatch(heartbeatLatch);
+			heartbeatLatch.await(5, TimeUnit.SECONDS);
+		} catch (InterruptedException e) {
+			logger.error("Error in testExpiredLocks", e);
+		}
+			//At this point the heartbeat object should hve
+			//refreshed the lock. assert that resource1 is
+			//locked
 		assertTrue(distLockFeat.beforeIsLocked("resource1"));
 	}