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() );
 }