From e190105e3533382e739ecafcf7678031d5441277 Mon Sep 17 00:00:00 2001
From: Felipe Gasparini <felipegasparini@gmail.com>
Date: Wed, 7 Feb 2018 23:59:36 +0100
Subject: [PATCH] Added unit tests for ILeaseManager injection in Worker
 Builder

---
 .../clientlibrary/lib/worker/Worker.java      |  7 +++++
 .../clientlibrary/lib/worker/WorkerTest.java  | 27 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java
index 3c39aa33..daf89c28 100644
--- a/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java
+++ b/src/main/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/Worker.java
@@ -426,6 +426,13 @@ public class Worker implements Runnable {
         return applicationName;
     }
 
+    /**
+     * @return the leaseCoordinator
+     */
+    KinesisClientLibLeaseCoordinator getLeaseCoordinator(){
+        return leaseCoordinator;
+    }
+
     /**
      * Start consuming data from the stream, and pass it to the application record processors.
      */
diff --git a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java
index 6cc7ef08..037a54b2 100644
--- a/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java
+++ b/src/test/java/com/amazonaws/services/kinesis/clientlibrary/lib/worker/WorkerTest.java
@@ -1500,6 +1500,33 @@ public class WorkerTest {
         Assert.assertTrue(worker.getStreamConfig().getStreamProxy() instanceof KinesisLocalFileProxy);
     }
 
+    @Test
+    public void testBuilderWithDefaultLeaseManager()  {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+
+        Worker worker = new Worker.Builder()
+                .recordProcessorFactory(recordProcessorFactory)
+                .config(config)
+                .build();
+
+        Assert.assertNotNull(worker.getLeaseCoordinator().getLeaseManager());
+    }
+
+    @SuppressWarnings("unchecked")
+    @Test
+    public void testBuilderWhenLeaseManagerIsSet()  {
+        IRecordProcessorFactory recordProcessorFactory = mock(IRecordProcessorFactory.class);
+        // Create an instance of ILeaseManager for injection and validation
+        ILeaseManager<KinesisClientLease> leaseManager = (ILeaseManager<KinesisClientLease>) mock(ILeaseManager.class);
+        Worker worker = new Worker.Builder()
+                .recordProcessorFactory(recordProcessorFactory)
+                .config(config)
+                .leaseManager(leaseManager)
+                .build();
+
+        Assert.assertSame(leaseManager, worker.getLeaseCoordinator().getLeaseManager());
+    }
+
     private abstract class InjectableWorker extends Worker {
         InjectableWorker(String applicationName, IRecordProcessorFactory recordProcessorFactory,
                 KinesisClientLibConfiguration config, StreamConfig streamConfig,