diff --git a/.gitignore b/.gitignore
index 90f483f..1ccc2a4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,13 +23,13 @@
 **/bal-api-3.1.0
 
 # Test binary
-test_openolt
-
-# proto files
-protos/voltha_protos/*.proto
+**/test_openolt
 
 # Test report
 **/test_openolt_report.xml
 
+# Downloaded voltha-protos file
+voltha_protos
+
 # IntelliJ Files
 .idea
diff --git a/agent/common/core.h b/agent/common/core.h
index 1caf3cc..26f1dac 100644
--- a/agent/common/core.h
+++ b/agent/common/core.h
@@ -107,9 +107,11 @@
 Status RemoveTrafficQueues_(const tech_profile::TrafficQueues *traffic_queues);
 uint32_t GetPortNum_(uint32_t flow_id);
 int get_status_bcm_cli_quit(void);
-uint16_t get_dev_id(void); 
+uint16_t get_dev_id(void);
 Status pushOltOperInd(uint32_t intf_id, const char *type, const char *state);
 uint64_t get_flow_status(uint16_t flow_id, uint16_t flow_type, uint16_t data_id);
 
 void stats_collection();
+Status check_connection();
+Status check_bal_ready();
 #endif
diff --git a/agent/scripts/init.d/openolt b/agent/scripts/init.d/openolt
index 1db4fbb..70594b1 100755
--- a/agent/scripts/init.d/openolt
+++ b/agent/scripts/init.d/openolt
@@ -25,9 +25,6 @@
 
 start() {
   printf "Starting '$NAME'... "
-  # TODO: If openolt agent has a way to figure out progamatically
-  # if the BAL is ready, the below timer may not be needed.
-  sleep 120
   rm -f /var/log/$NAME.log
   touch /var/log/$NAME.log
   start-stop-daemon --verbose --start --chuid "$USER:$GROUP" --background --no-close --make-pidfile --pidfile /var/run/$NAME.pid --chdir "$APPDIR" --exec "$APPBIN" -- $APPARGS < /dev/tty1 >> /var/log/$NAME.log 2>&1 || true
diff --git a/agent/src/core.cc b/agent/src/core.cc
index ef92547..897d326 100644
--- a/agent/src/core.cc
+++ b/agent/src/core.cc
@@ -25,6 +25,7 @@
 #include <thread>
 #include <bitset>
 #include <inttypes.h>
+#include <unistd.h>
 
 #include "device.h"
 #include "core.h"
@@ -161,6 +162,9 @@
     return dev_id;
 }
 
+// Stubbed defntion of bcmolt_cfg_get required for unit-test
+extern bcmos_errno bcmolt_cfg_get_stub(bcmolt_oltid olt_id, void* ptr);
+
 /**
 * Returns the default NNI (Upstream direction) or PON (Downstream direction) scheduler
 * Every NNI port and PON port have default scheduler.
@@ -782,8 +786,11 @@
         bcmos_fastlock_init(&data_lock, 0);
         OPENOLT_LOG(INFO, openolt_log_id, "Enable OLT - %s-%s\n", VENDOR_ID, MODEL_ID);
 
-        if (bcmolt_api_conn_mgr_is_connected(dev_id))
-        {
+        //check BCM daemon is connected or not
+        Status status = check_connection();
+        if (!status.ok())
+            return status;
+        else {
             Status status = SubscribeIndication();
             if (!status.ok()) {
                 OPENOLT_LOG(ERROR, openolt_log_id, "SubscribeIndication failed - %s : %s\n",
@@ -791,6 +798,14 @@
                     status.error_message().c_str());
                 return status;
             }
+
+            //check BAL state in initial stage
+            status = check_bal_ready();
+            if (!status.ok())
+                return status;
+        }
+
+        {
             bcmos_errno err;
             bcmolt_odid dev;
             OPENOLT_LOG(INFO, openolt_log_id, "Enabling PON %d Devices ... \n", BCM_MAX_DEVS_PER_LINE_CARD);
@@ -801,7 +816,7 @@
                 BCMOLT_CFG_INIT(&dev_cfg, device, dev_key);
                 BCMOLT_MSG_FIELD_GET(&dev_cfg, system_mode);
                 err = bcmolt_cfg_get(dev_id, &dev_cfg.hdr);
-                if (err == BCM_ERR_NOT_CONNECTED) { 
+                if (err == BCM_ERR_NOT_CONNECTED) {
                     bcmolt_device_key key = {.device_id = dev};
                     bcmolt_device_connect oper;
                     BCMOLT_OPER_INIT(&oper, device, connect, key);
@@ -1148,7 +1163,7 @@
     err = get_pon_interface_status((bcmolt_interface)intf_id, &state); 
     if (err == BCM_ERR_OK) {
         if (state == BCMOLT_INTERFACE_STATE_ACTIVE_WORKING) {
-            OPENOLT_LOG(INFO, openolt_log_id, "PON interface: %d already enabled\n", intf_id);
+            OPENOLT_LOG(WARNING, openolt_log_id, "PON interface: %d already enabled\n", intf_id);
             return Status::OK;
         }
     } 
@@ -1347,7 +1362,7 @@
     err = get_nni_interface_status((bcmolt_interface)intf_id, &state); 
     if (err == BCM_ERR_OK) {
         if (set_state && state == BCMOLT_INTERFACE_STATE_ACTIVE_WORKING) {
-            OPENOLT_LOG(INFO, openolt_log_id, "NNI interface: %d already enabled\n", intf_id);
+            OPENOLT_LOG(WARNING, openolt_log_id, "NNI interface: %d already enabled\n", intf_id);
             OPENOLT_LOG(INFO, openolt_log_id, "Initializing tm sched creation for NNI interface: %d\n", intf_id);
             CreateDefaultSched(intf_id, upstream);
             CreateDefaultQueue(intf_id, upstream);
@@ -2190,6 +2205,19 @@
     bcmolt_tm_sched_key tm_sched_key = {.id = 1};
     tm_sched_key.id = get_default_tm_sched_id(intf_id, direction);
 
+    //check TM scheduler has configured or not
+    BCMOLT_CFG_INIT(&tm_sched_cfg, tm_sched, tm_sched_key);
+    BCMOLT_MSG_FIELD_GET(&tm_sched_cfg, state);
+    err = bcmolt_cfg_get(dev_id, &tm_sched_cfg.hdr);
+    if (err) {
+        OPENOLT_LOG(ERROR, openolt_log_id, "cfg: Failed to query TM scheduler\n");
+        return err;
+    }
+    else if (tm_sched_cfg.data.state == BCMOLT_CONFIG_STATE_CONFIGURED) {
+        OPENOLT_LOG(WARNING, openolt_log_id, "tm scheduler default config has already with id %d\n", tm_sched_key.id);
+        return BCM_ERR_OK;
+    }
+
     // bcmbal_tm_sched_owner
     BCMOLT_CFG_INIT(&tm_sched_cfg, tm_sched, tm_sched_key);
 
@@ -2799,3 +2827,49 @@
     }
     return Status::OK;
 }
+
+Status check_connection() {
+    int maxTrials = 60;
+    while (!bcmolt_api_conn_mgr_is_connected(dev_id)) {
+        sleep(1);
+        if (--maxTrials == 0)
+            return grpc::Status(grpc::StatusCode::UNAVAILABLE, "check connection failed");
+        else
+            OPENOLT_LOG(INFO, openolt_log_id, "waiting for daemon connection ...\n");
+    }
+    OPENOLT_LOG(INFO, openolt_log_id, "daemon is connected\n");
+    return Status::OK;
+}
+
+Status check_bal_ready() {
+    bcmos_errno err;
+    int maxTrials = 30;
+    bcmolt_olt_cfg olt_cfg = { };
+    bcmolt_olt_key olt_key = { };
+
+    BCMOLT_CFG_INIT(&olt_cfg, olt, olt_key);
+    BCMOLT_MSG_FIELD_GET(&olt_cfg, bal_state);
+
+    while (olt_cfg.data.bal_state != BCMOLT_BAL_STATE_BAL_AND_SWITCH_READY) {
+        if (--maxTrials == 0)
+            return grpc::Status(grpc::StatusCode::UNAVAILABLE, "check bal ready failed");
+        sleep(5);
+        #ifdef TEST_MODE
+        // It is impossible to mock the setting of olt_cfg.data.bal_state because
+        // the actual bcmolt_cfg_get passes the address of olt_cfg.hdr and we cannot
+        // set the olt_cfg.data.bal_state. So a new stub function is created and address
+        // of olt_cfg is passed. This is one-of case where we need to add specific
+        // code in product code.
+        if (bcmolt_cfg_get_stub(dev_id, &olt_cfg)) {
+        #else
+        if (bcmolt_cfg_get(dev_id, &olt_cfg.hdr)) {
+        #endif
+            continue;
+        }
+        else
+            OPENOLT_LOG(INFO, openolt_log_id, "waiting for BAL ready ...\n");
+    }
+
+    OPENOLT_LOG(INFO, openolt_log_id, "BAL is ready\n");
+    return Status::OK;
+}
diff --git a/agent/src/indications.cc b/agent/src/indications.cc
index 1826b77..cc42adc 100644
--- a/agent/src/indications.cc
+++ b/agent/src/indications.cc
@@ -145,7 +145,6 @@
 static void OltOperIndication(bcmolt_devid olt, bcmolt_msg *msg) {
     openolt::Indication ind;
     openolt::OltIndication* olt_ind = new openolt::OltIndication;
-    Status status;
     std::string admin_state;
 
     switch (msg->subgroup) {
diff --git a/agent/test/Makefile.in b/agent/test/Makefile.in
index 2122266..3f8eaf9 100644
--- a/agent/test/Makefile.in
+++ b/agent/test/Makefile.in
@@ -79,7 +79,7 @@
 	sudo make -C $(GRPC_DST) install
 	sudo ldconfig
 
-prereq-mock-lib: gtest cmock
+prereq-mock-lib: gtest cmock gmock-global
 
 gtest:
 	# Install gtest and gmock
@@ -96,6 +96,11 @@
 	make -C /tmp/C-Mock
 	sudo make -C /tmp/C-Mock install
 
+gmock-global:
+	# Install gmock-global
+	rm -rf /tmp/gmock-global && cd /tmp && git clone https://github.com/apriorit/gmock-global.git
+	cd /tmp/gmock-global && git checkout 1.0.2
+	sudo cp -rf /tmp/gmock-global/include/gmock-global /usr/local/include/
 ########################################################################
 ##
 ##
diff --git a/agent/test/inc/bal_mocker.h b/agent/test/inc/bal_mocker.h
index 85bdb22..72cbbc4 100644
--- a/agent/test/inc/bal_mocker.h
+++ b/agent/test/inc/bal_mocker.h
@@ -18,7 +18,10 @@
 #ifndef __BAL_MOCKER_H__
 #define __BAL_MOCKER_H__
 #include <cmock/cmock.h>
+#include <gmock-global/gmock-global.h>
 #include <cstdlib>
+#include <grpc++/grpc++.h>
+using grpc::Status;
 extern "C" {
 #include <bcmos_system.h>
 #include <bcmolt_api.h>
@@ -41,5 +44,7 @@
     MOCK_METHOD1(bcmolt_host_init, bcmos_errno(bcmolt_host_init_parms*));
     MOCK_METHOD2(bcmolt_cfg_get, bcmos_errno(bcmolt_oltid, bcmolt_cfg*));
     MOCK_METHOD2(bcmolt_oper_submit, bcmos_errno(bcmolt_oltid, bcmolt_oper*));
+  // Add more here
 };
+
 #endif
diff --git a/agent/test/src/bal_mocker.cc b/agent/test/src/bal_mocker.cc
index b38906f..1f9d4b0 100644
--- a/agent/test/src/bal_mocker.cc
+++ b/agent/test/src/bal_mocker.cc
@@ -23,3 +23,4 @@
 CMOCK_MOCK_FUNCTION2(BalMocker, bcmolt_cfg_get, bcmos_errno(bcmolt_oltid, bcmolt_cfg*));
 CMOCK_MOCK_FUNCTION2(BalMocker, bcmolt_oper_submit, bcmos_errno(bcmolt_oltid, bcmolt_oper*));
 }
+
diff --git a/agent/test/src/main_test.cc b/agent/test/src/main_test.cc
index 7f04e9c..8e3a903 100644
--- a/agent/test/src/main_test.cc
+++ b/agent/test/src/main_test.cc
@@ -17,7 +17,9 @@
 
 #include "gtest/gtest.h"
 
+using namespace testing;
+
 int main(int argc, char **argv) {
-  ::testing::InitGoogleTest(&argc, argv);
+  InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();
 }
diff --git a/agent/test/src/test_enable_olt.cc b/agent/test/src/test_enable_olt.cc
index 0cdb965..e7757ca 100644
--- a/agent/test/src/test_enable_olt.cc
+++ b/agent/test/src/test_enable_olt.cc
@@ -18,7 +18,9 @@
 #include "bal_mocker.h"
 #include "core.h"
 
-class TestOltEnable : public ::testing::Test {
+using namespace testing;
+
+class TestOltEnable : public Test {
  protected:
   virtual void SetUp() {
   }
@@ -29,28 +31,41 @@
   }
 };
 
+// This is used to set custom bcmolt_cfg value to bcmolt_cfg pointer coming in
+// bcmolt_cfg_get_stub.
+ACTION_P(SetArg1ToBcmOltCfg, value) { *static_cast<bcmolt_olt_cfg*>(arg1) = value; };
+
+
+// Create a mock function for bcmolt_cfg_get_stub C++ function
+MOCK_GLOBAL_FUNC2(bcmolt_cfg_get_stub, bcmos_errno(bcmolt_oltid, void*));
+
 
 // Test Fixture for OltEnable
 
 // Test 1: OltEnableSuccess case
 TEST_F(TestOltEnable, OltEnableSuccess){
-    // NiceMock is used to suppress 'WillByDefault' return errors on using 'NiceMock'.
+    // NiceMock is used to suppress 'WillByDefault' return errors.
     // This is described in https://github.com/arangodb-helper/gtest/blob/master/googlemock/docs/CookBook.md
-    ::testing::NiceMock<BalMocker> balMock;
+    NiceMock<BalMocker> balMock;
     bcmos_errno host_init_res = BCM_ERR_OK;
+    bcmos_errno bal_cfg_get_stub_res = BCM_ERR_OK;
     bcmos_errno bal_cfg_get_res = BCM_ERR_NOT_CONNECTED;
     bcmos_errno olt_oper_res = BCM_ERR_OK;
 
+    bcmolt_olt_cfg olt_cfg = { };
+    bcmolt_olt_key olt_key = { };
+    BCMOLT_CFG_INIT(&olt_cfg, olt, olt_key);
+    olt_cfg.data.bal_state = BCMOLT_BAL_STATE_BAL_AND_SWITCH_READY;
+
     Status olt_enable_res;
 
-    // The 'EXPECT_CALL' will do strict validation of input parameters. This may not be relevant for
-    // the current test case. Use 'ON_CALL' instead.
-    // The ON_CALL results in WARNINGs when running tests. Use NickMock instead of directly using BalMocker.
-    // https://github.com/google/googletest/blob/master/googlemock/docs/cook_book.md#the-nice-the-strict-and-the-naggy-nicestrictnaggy
-    // In below tests '::testing::_' does no validation on argument.
-    ON_CALL(balMock, bcmolt_host_init(::testing::_)).WillByDefault(::testing::Return(host_init_res));
-    ON_CALL(balMock, bcmolt_cfg_get(::testing::_,::testing::_)).WillByDefault(::testing::Return(bal_cfg_get_res));
-    ON_CALL(balMock, bcmolt_oper_submit(::testing::_, ::testing::_)).WillByDefault(::testing::Return(olt_oper_res));
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get_stub, bcmolt_cfg_get_stub(_, _))
+                     .WillOnce(DoAll(SetArg1ToBcmOltCfg(olt_cfg), Return(bal_cfg_get_stub_res)));
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(BCM_MAX_DEVS_PER_LINE_CARD)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
+    ON_CALL(balMock, bcmolt_oper_submit(_, _)).WillByDefault(Return(olt_oper_res));
 
     olt_enable_res = Enable_(1, NULL);
     ASSERT_TRUE( olt_enable_res.error_message() == Status::OK.error_message() );
@@ -58,23 +73,17 @@
 
 // Test 2: OltEnableFail_host_init_fail
 TEST_F(TestOltEnable, OltEnableFail_host_init_fail) {
-    // NiceMock is used to suppress 'WillByDefault' return errors on using 'NiceMock'.
+    // NiceMock is used to suppress 'WillByDefault' return errors.
     // This is described in https://github.com/arangodb-helper/gtest/blob/master/googlemock/docs/CookBook.md
-    ::testing::NiceMock<BalMocker> balMock;
+    NiceMock<BalMocker> balMock;
     bcmos_errno host_init_res = BCM_ERR_INTERNAL;
-    bcmos_errno bal_cfg_get_res = BCM_ERR_NOT_CONNECTED;
-    bcmos_errno olt_oper_res = BCM_ERR_OK;
 
     Status olt_enable_res;
 
-    // The 'EXPECT_CALL' will do strict validation of input parameters. This may not be relevant for
-    // the current test case. Use 'ON_CALL' instead.
-    // The ON_CALL results in WARNINGs when running tests. Use NickMock instead of directly using BalMocker.
-    // https://github.com/google/googletest/blob/master/googlemock/docs/cook_book.md#the-nice-the-strict-and-the-naggy-nicestrictnaggy
-    // In below tests '::testing::_' does no validation on argument.
-    ON_CALL(balMock, bcmolt_host_init(::testing::_)).WillByDefault(::testing::Return(host_init_res));
-    ON_CALL(balMock, bcmolt_cfg_get(::testing::_,::testing::_)).WillByDefault(::testing::Return(bal_cfg_get_res));
-    ON_CALL(balMock, bcmolt_oper_submit(::testing::_, ::testing::_)).WillByDefault(::testing::Return(olt_oper_res));
+    // Ensure that the state of the OLT is in deactivated to start with..
+    state.deactivate();
+
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
 
     olt_enable_res = Enable_(1, NULL);
     ASSERT_TRUE( olt_enable_res.error_message() != Status::OK.error_message() );
@@ -82,81 +91,104 @@
 
 // Test 3: OltEnableSuccess_PON_Device_Connected
 TEST_F(TestOltEnable, OltEnableSuccess_PON_Device_Connected) {
-    // NiceMock is used to suppress 'WillByDefault' return errors on using 'NiceMock'.
+
+    // NiceMock is used to suppress 'WillByDefault' return errors.
     // This is described in https://github.com/arangodb-helper/gtest/blob/master/googlemock/docs/CookBook.md
-    ::testing::NiceMock<BalMocker> balMock;
+    NiceMock<BalMocker> balMock;
     bcmos_errno host_init_res = BCM_ERR_OK;
+    bcmos_errno bal_cfg_get_stub_res = BCM_ERR_OK;
     bcmos_errno bal_cfg_get_res = BCM_ERR_OK;
+    bcmos_errno olt_oper_res = BCM_ERR_OK;
 
-    Status olt_enable_res;
-
-    // The 'EXPECT_CALL' will do strict validation of input parameters. This may not be relevant for
-    // the current test case. Use 'ON_CALL' instead.
-    // The ON_CALL results in WARNINGs when running tests. Use NickMock instead of directly using BalMocker.
-    // https://github.com/google/googletest/blob/master/googlemock/docs/cook_book.md#the-nice-the-strict-and-the-naggy-nicestrictnaggy
-    // In below tests '::testing::_' does no validation on argument.
-    ON_CALL(balMock, bcmolt_host_init(::testing::_)).WillByDefault(::testing::Return(host_init_res));
-    ON_CALL(balMock, bcmolt_cfg_get(::testing::_,::testing::_)).WillByDefault(::testing::Return(bal_cfg_get_res));
-
-    olt_enable_res = Enable_(1, NULL);
-    ASSERT_TRUE( olt_enable_res.error_message() == Status::OK.error_message() );
-}
-
-// Test 4: OltEnableFail_All_PON_Enable_Fail
-TEST_F(TestOltEnable, OltEnableFail_All_PON_Enable_Fail) {
-    // NiceMock is used to suppress 'WillByDefault' return errors on using 'NiceMock'.
-    // This is described in https://github.com/arangodb-helper/gtest/blob/master/googlemock/docs/CookBook.md
-    ::testing::NiceMock<BalMocker> balMock;
-    bcmos_errno host_init_res = BCM_ERR_OK;
-    bcmos_errno bal_cfg_get_res = BCM_ERR_NOT_CONNECTED;
-    bcmos_errno olt_oper_res = BCM_ERR_INTERNAL;
+    bcmolt_olt_cfg olt_cfg = { };
+    bcmolt_olt_key olt_key = { };
+    BCMOLT_CFG_INIT(&olt_cfg, olt, olt_key);
+    olt_cfg.data.bal_state = BCMOLT_BAL_STATE_BAL_AND_SWITCH_READY;
 
     Status olt_enable_res;
 
     // Ensure that the state of the OLT is in deactivated to start with..
     state.deactivate();
 
-    // The 'EXPECT_CALL' will do strict validation of input parameters. This may not be relevant for
-    // the current test case. Use 'ON_CALL' instead.
-    // The ON_CALL results in WARNINGs when running tests. Use NickMock instead of directly using BalMocker.
-    // https://github.com/google/googletest/blob/master/googlemock/docs/cook_book.md#the-nice-the-strict-and-the-naggy-nicestrictnaggy
-    // In below tests '::testing::_' does no validation on argument.
-    ON_CALL(balMock, bcmolt_host_init(::testing::_)).WillByDefault(::testing::Return(host_init_res));
-    ON_CALL(balMock, bcmolt_cfg_get(::testing::_,::testing::_)).WillByDefault(::testing::Return(bal_cfg_get_res));
-    ON_CALL(balMock, bcmolt_oper_submit(::testing::_, ::testing::_)).WillByDefault(::testing::Return(olt_oper_res));
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get_stub, bcmolt_cfg_get_stub(_, _))
+                     .WillOnce(DoAll(SetArg1ToBcmOltCfg(olt_cfg), Return(bal_cfg_get_stub_res)));
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(BCM_MAX_DEVS_PER_LINE_CARD)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
 
     olt_enable_res = Enable_(1, NULL);
+    ASSERT_TRUE( olt_enable_res.error_message() == Status::OK.error_message() );
+
+}
+
+// Test 4: OltEnableFail_All_PON_Enable_Fail
+TEST_F(TestOltEnable, OltEnableFail_All_PON_Enable_Fail) {
+
+    // NiceMock is used to suppress 'WillByDefault' return errors.
+    // This is described in https://github.com/arangodb-helper/gtest/blob/master/googlemock/docs/CookBook.md
+    NiceMock<BalMocker> balMock;
+    bcmos_errno host_init_res = BCM_ERR_OK;
+    bcmos_errno bal_cfg_get_stub_res = BCM_ERR_OK;
+    bcmos_errno bal_cfg_get_res = BCM_ERR_NOT_CONNECTED;
+    bcmos_errno olt_oper_res = BCM_ERR_INTERNAL;
+
+    bcmolt_olt_cfg olt_cfg = { };
+    bcmolt_olt_key olt_key = { };
+    BCMOLT_CFG_INIT(&olt_cfg, olt, olt_key);
+    olt_cfg.data.bal_state = BCMOLT_BAL_STATE_BAL_AND_SWITCH_READY;
+
+    Status olt_enable_res;
+
+    // Ensure that the state of the OLT is in deactivated to start with..
+    state.deactivate();
+
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get_stub, bcmolt_cfg_get_stub(_, _))
+                     .WillOnce(DoAll(SetArg1ToBcmOltCfg(olt_cfg), Return(bal_cfg_get_stub_res)));
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(BCM_MAX_DEVS_PER_LINE_CARD)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
+    ON_CALL(balMock, bcmolt_oper_submit(_, _)).WillByDefault(Return(olt_oper_res));
+
+    olt_enable_res = Enable_(1, NULL);
+
     ASSERT_TRUE( olt_enable_res.error_message() != Status::OK.error_message() );
 }
 
 // Test 5 OltEnableSuccess_One_PON_Enable_Fail : One PON device enable fails, but all others succeed.
 TEST_F(TestOltEnable, OltEnableSuccess_One_PON_Enable_Fail) {
-    // NiceMock is used to suppress 'WillByDefault' return errors on using 'NiceMock'.
+
+    // NiceMock is used to suppress 'WillByDefault' return errors.
     // This is described in https://github.com/arangodb-helper/gtest/blob/master/googlemock/docs/CookBook.md
-    ::testing::NiceMock<BalMocker> balMock;
+    NiceMock<BalMocker> balMock;
     bcmos_errno host_init_res = BCM_ERR_OK;
+    bcmos_errno bal_cfg_get_stub_res = BCM_ERR_OK;
     bcmos_errno bal_cfg_get_res = BCM_ERR_NOT_CONNECTED;
     bcmos_errno olt_oper_res_fail = BCM_ERR_INTERNAL;
     bcmos_errno olt_oper_res_success = BCM_ERR_OK;
 
+    bcmolt_olt_cfg olt_cfg = { };
+    bcmolt_olt_key olt_key = { };
+    BCMOLT_CFG_INIT(&olt_cfg, olt, olt_key);
+    olt_cfg.data.bal_state = BCMOLT_BAL_STATE_BAL_AND_SWITCH_READY;
+
     Status olt_enable_res;
 
     // Ensure that the state of the OLT is in deactivated to start with..
     state.deactivate();
 
-    // The 'EXPECT_CALL' will do strict validation of input parameters. This may not be relevant for
-    // the current test case. Use 'ON_CALL' instead.
-    // The ON_CALL results in WARNINGs when running tests. Use NickMock instead of directly using BalMocker.
-    // https://github.com/google/googletest/blob/master/googlemock/docs/cook_book.md#the-nice-the-strict-and-the-naggy-nicestrictnaggy
-    // In below tests '::testing::_' does no validation on argument.
-    ON_CALL(balMock, bcmolt_host_init(::testing::_)).WillByDefault(::testing::Return(host_init_res));
-    ON_CALL(balMock, bcmolt_cfg_get(::testing::_,::testing::_)).WillByDefault(::testing::Return(bal_cfg_get_res));
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get_stub, bcmolt_cfg_get_stub(_, _))
+                     .WillOnce(DoAll(SetArg1ToBcmOltCfg(olt_cfg), Return(bal_cfg_get_stub_res)));
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(BCM_MAX_DEVS_PER_LINE_CARD)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
     // For the the first PON mac device, the activation result will fail, and will succeed for all other PON mac devices.
-    EXPECT_CALL(balMock, bcmolt_oper_submit(::testing::_, ::testing::_))
-                         .WillOnce(::testing::Return(olt_oper_res_fail))
-                         .WillRepeatedly(::testing::Return(olt_oper_res_success));
-
+    EXPECT_CALL(balMock, bcmolt_oper_submit(_, _))
+                         .WillOnce(Return(olt_oper_res_fail))
+                         .WillRepeatedly(Return(olt_oper_res_success));
     olt_enable_res = Enable_(1, NULL);
-    // The OLT activation should succeed if at least one PON mac device activation succeeds.
+
     ASSERT_TRUE( olt_enable_res.error_message() == Status::OK.error_message() );
 }