diff --git a/agent/common/main.cc b/agent/common/main.cc
index a18db3c..49d1175 100644
--- a/agent/common/main.cc
+++ b/agent/common/main.cc
@@ -99,7 +99,12 @@
         }
     }
 
-    ProbeDeviceCapabilities_();
+    status = ProbeDeviceCapabilities_();
+    if (!status.ok()) {
+        std::cout << "ERROR: Could not find the OLT Device capabilities" << std::endl;
+        return 1;
+    }
+
     sleep(2);
     // Enable all PON interfaces. 
     for (int i = 0; i < NumPonIf_(); i++) {
diff --git a/agent/src/core.cc b/agent/src/core.cc
index 8de776f..a3653f3 100644
--- a/agent/src/core.cc
+++ b/agent/src/core.cc
@@ -103,6 +103,10 @@
 const std::string downstream = "downstream";
 bcmolt_oltid dev_id = 0;
 
+/* Constants used for retrying some BAL APIs */
+const uint32_t BAL_API_RETRY_TIME_IN_USECS = 1000000;
+const uint32_t MAX_BAL_API_RETRY_COUNT = 5;
+
 /* Current session */
 static bcmcli_session *current_session;
 static bcmcli_entry *api_parent_dir;
@@ -1299,6 +1303,36 @@
     return Status::OK;
 }
 
+/* Same as bcmolt_cfg_get but with added logic of retrying the API
+   in case of some specific failures like timeout or object not yet ready
+*/
+bcmos_errno bcmolt_cfg_get_mult_retry(bcmolt_oltid olt, bcmolt_cfg *cfg) {
+    bcmos_errno err;
+    uint32_t current_try = 0;
+
+    while (current_try < MAX_BAL_API_RETRY_COUNT) {
+        err = bcmolt_cfg_get(olt, cfg);
+        current_try++;
+
+        if (err == BCM_ERR_STATE || err == BCM_ERR_TIMEOUT) {
+            OPENOLT_LOG(WARNING, openolt_log_id, "bcmolt_cfg_get: err = %s(%d)\n",bcmos_strerror(err), err);
+            bcmos_usleep(BAL_API_RETRY_TIME_IN_USECS);
+            continue;
+        }
+        else {
+           break;
+        }
+    }
+
+    if (err != BCM_ERR_OK) {
+        OPENOLT_LOG(ERROR, openolt_log_id, "bcmolt_cfg_get tried (%d) times with retry time(%d usecs) err=%d\n",
+                           current_try,
+                           BAL_API_RETRY_TIME_IN_USECS,
+                           err);
+    }
+    return err;
+}
+
 Status ProbeDeviceCapabilities_() {
     bcmos_errno err;
     bcmolt_device_cfg dev_cfg = { };
@@ -1323,11 +1357,11 @@
         // code in production code.
     err = bcmolt_cfg_get__olt_topology_stub(dev_id, &olt_cfg);
     #else
-    err = bcmolt_cfg_get(dev_id, &olt_cfg.hdr);
+    err = bcmolt_cfg_get_mult_retry(dev_id, &olt_cfg.hdr);
     #endif
     if (err) {
-        OPENOLT_LOG(ERROR, openolt_log_id, "cfg: Failed to query OLT\n");
-        return bcm_to_grpc_err(err, "cfg: Failed to query OLT");
+        OPENOLT_LOG(ERROR, openolt_log_id, "cfg: Failed to query OLT topology\n");
+        return bcm_to_grpc_err(err, "cfg: Failed to query OLT topology");
     }
 
     num_of_nni_ports = olt_cfg.data.topology.num_switch_ports;
@@ -1342,16 +1376,18 @@
             num_of_pon_ports,
             BCM_MAX_DEVS_PER_LINE_CARD);
 
+    uint32_t num_failed_cfg_gets = 0;
     for (int devid = 0; devid < BCM_MAX_DEVS_PER_LINE_CARD; devid++) {
         dev_key.device_id = devid;
         BCMOLT_CFG_INIT(&dev_cfg, device, dev_key);
         BCMOLT_MSG_FIELD_GET(&dev_cfg, firmware_sw_version);
         BCMOLT_MSG_FIELD_GET(&dev_cfg, chip_family);
         BCMOLT_MSG_FIELD_GET(&dev_cfg, system_mode);
-        err = bcmolt_cfg_get(dev_id, &dev_cfg.hdr);
+        err = bcmolt_cfg_get_mult_retry(dev_id, &dev_cfg.hdr);
         if (err) {
-            OPENOLT_LOG(ERROR, openolt_log_id, "device: Failed to query OLT\n");
-            return bcm_to_grpc_err(err, "device: Failed to query OLT");
+            OPENOLT_LOG(WARNING, openolt_log_id,"Failed to query PON MAC Device %d (errno = %d). Skipping the device.\n", devid, err);
+            num_failed_cfg_gets++;
+            continue;
         }
 
         std::string bal_version;
@@ -1384,6 +1420,13 @@
         bcmos_usleep(500000);
     }
 
+    /* If all the devices returned errors then we tell the caller that this is an error else we work with 
+       only the devices that retured success*/
+    if (num_failed_cfg_gets == BCM_MAX_DEVS_PER_LINE_CARD) {
+        OPENOLT_LOG(ERROR, openolt_log_id, "device: Query of all the devices failed\n");
+        return bcm_to_grpc_err(err, "device: All devices failed query");
+    }
+
     return Status::OK;
 }
 #if 0
diff --git a/agent/test/Makefile.in b/agent/test/Makefile.in
index b08f3c5..6351622 100644
--- a/agent/test/Makefile.in
+++ b/agent/test/Makefile.in
@@ -118,7 +118,7 @@
 TEST_BIN=test_openolt
 # The below way of getting source files is not working
 #TEST_SRCS = $($(filter-out src/bal_stub.cc, $(wildcard src/*.cc)))
-TEST_SRCS = src/main_test.cc src/bal_mocker.cc src/test_enable_olt.cc src/test_disable_reenable_olt.cc
+TEST_SRCS = src/main_test.cc src/bal_mocker.cc src/test_core.cc 
 TEST_OBJS = $(TEST_SRCS:.cc=.o)
 build: $(TEST_OBJS)
 	$(CXX) $(shell cmock-config --libs) -L$(TEST_OPENOLT_LIB_DIR) -L/usr/lib/ -I/usr/include -I./inc -o $(TEST_BIN) $(TEST_OBJS) -lgmock -lgtest -lopenolt -lopenoltapi  $(LDFLAGS)
diff --git a/agent/test/src/bal_stub.cc b/agent/test/src/bal_stub.cc
index 6d757b3..f3f27b4 100644
--- a/agent/test/src/bal_stub.cc
+++ b/agent/test/src/bal_stub.cc
@@ -31,8 +31,8 @@
 dev_log_id def_log_id=0;
 
 void bcmos_usleep(uint32_t us) {
-    // let always sleep for 1s irrespective of the value passed.
-    sleep (1);
+    // let always sleep for 10ms irrespective of the value passed.
+    usleep (10000);
 }
 
 void bcmos_fastlock_init(bcmos_fastlock *lock, uint32_t flags)  {
diff --git a/agent/test/src/test_core.cc b/agent/test/src/test_core.cc
new file mode 100644
index 0000000..3588f95
--- /dev/null
+++ b/agent/test/src/test_core.cc
@@ -0,0 +1,386 @@
+/*
+ * Copyright 2018-present Open Networking Foundation
+
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+
+ * http://www.apache.org/licenses/LICENSE-2.0
+
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "gtest/gtest.h"
+#include "bal_mocker.h"
+#include "core.h"
+using namespace testing;
+using namespace std;
+
+class TestOltEnable : public Test {
+ protected:
+  virtual void SetUp() {
+  }
+
+  virtual void TearDown() {
+    // Code here will be called immediately after each test
+    // (right before the destructor).
+  }
+};
+
+// This is used to set custom bcmolt_cfg value to bcmolt_cfg pointer coming in
+// bcmolt_cfg_get__bal_state_stub.
+ACTION_P(SetArg1ToBcmOltCfg, value) { *static_cast<bcmolt_olt_cfg*>(arg1) = value; };
+
+
+// Create a mock function for bcmolt_cfg_get__bal_state_stub C++ function
+MOCK_GLOBAL_FUNC2(bcmolt_cfg_get__bal_state_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.
+    // 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_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;
+
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__bal_state_stub, bcmolt_cfg_get__bal_state_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 2: OltEnableFail_host_init_fail
+TEST_F(TestOltEnable, OltEnableFail_host_init_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_INTERNAL;
+
+    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));
+
+    olt_enable_res = Enable_(1, NULL);
+    ASSERT_TRUE( olt_enable_res.error_message() != Status::OK.error_message() );
+}
+
+// Test 3: OltEnableSuccess_PON_Device_Connected
+TEST_F(TestOltEnable, OltEnableSuccess_PON_Device_Connected) {
+
+    // 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_OK;
+    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;
+
+    // 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__bal_state_stub, bcmolt_cfg_get__bal_state_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__bal_state_stub, bcmolt_cfg_get__bal_state_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.
+    // 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_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();
+
+    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__bal_state_stub, bcmolt_cfg_get__bal_state_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(_, _))
+                         .WillOnce(Return(olt_oper_res_fail))
+                         .WillRepeatedly(Return(olt_oper_res_success));
+    olt_enable_res = Enable_(1, NULL);
+
+    ASSERT_TRUE( olt_enable_res.error_message() == Status::OK.error_message() );
+}
+
+////////////////////////////////////////////////////////////////////////
+// For testing Enable/Disable functionality
+////////////////////////////////////////////////////////////////////////
+
+int num_of_pon_ports = 16;
+
+// Create a mock function for bcmolt_cfg_get__olt_topology_stub C++ function
+MOCK_GLOBAL_FUNC2(bcmolt_cfg_get__olt_topology_stub, bcmos_errno(bcmolt_oltid, void*));
+
+class TestOltDisableReenable : public Test {
+ protected:
+  virtual void SetUp() {
+    NiceMock<BalMocker> balMock;
+    bcmos_errno bal_cfg_get_stub_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.topology.topology_maps.len = num_of_pon_ports;
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__olt_topology_stub, bcmolt_cfg_get__olt_topology_stub(_, _))
+                     .WillOnce(DoAll(SetArg1ToBcmOltCfg(olt_cfg), Return(bal_cfg_get_stub_res)));
+
+    ProbeDeviceCapabilities_();
+
+  }
+
+  virtual void TearDown() {
+    // Code here will be called immediately after each test
+    // (right before the destructor).
+  }
+};
+
+
+// Test Fixture for OltDisable
+
+// Test 1: OltDisableSuccess  case
+TEST_F(TestOltDisableReenable, OltDisableSuccess){
+    // 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 bal_cfg_get_res = BCM_ERR_OK;
+
+    Status olt_disable_res;
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(num_of_pon_ports)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
+
+    olt_disable_res = Disable_();
+    ASSERT_TRUE( olt_disable_res.error_message() == Status::OK.error_message() );
+
+}
+
+// Test 2: OltDisableAllPonFailed  case
+TEST_F(TestOltDisableReenable, OltDisableAllPonFailed){
+    // 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 bal_cfg_get_res = BCM_ERR_INTERNAL;
+
+    Status olt_disable_res;
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(num_of_pon_ports)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
+
+    olt_disable_res = Disable_();
+    ASSERT_TRUE( olt_disable_res.error_code() == grpc::StatusCode::INTERNAL);
+}
+
+
+// Test Fixture for OltReenable
+
+// Test 1: OltReenableSuccess  case
+TEST_F(TestOltDisableReenable, OltReenableSuccess){
+    // 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 bal_cfg_get_res = BCM_ERR_OK;
+
+    Status olt_reenable_res;
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(num_of_pon_ports*2)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
+
+    olt_reenable_res = Reenable_();
+    ASSERT_TRUE( olt_reenable_res.error_message() == Status::OK.error_message() );
+
+}
+
+// Test 2: OltReenableAllPonFailed  case
+TEST_F(TestOltDisableReenable, OltReenableAllPonFailed){
+    // 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 bal_cfg_get_res = BCM_ERR_OK;
+    bcmos_errno olt_oper_res = BCM_ERR_INTERNAL;
+
+    Status olt_reenable_res;
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+                         .Times(num_of_pon_ports)
+                         .WillRepeatedly(Return(bal_cfg_get_res));
+    EXPECT_CALL(balMock,bcmolt_oper_submit(_, _))
+                         .Times(num_of_pon_ports)
+                         .WillRepeatedly(Return(olt_oper_res));
+    olt_reenable_res = Reenable_();
+    ASSERT_TRUE( olt_reenable_res.error_code() == grpc::StatusCode::INTERNAL);
+}
+
+////////////////////////////////////////////////////////////////////////////
+// For testing ProbeDeviceCapabilities functionality
+////////////////////////////////////////////////////////////////////////////
+class TestProbeDevCapabilities : public Test {
+    protected:
+        NiceMock<BalMocker> balMock;
+        bcmos_errno olt_res_success = BCM_ERR_OK;
+        bcmos_errno olt_res_fail = BCM_ERR_COMM_FAIL;
+        bcmos_errno dev_res_success = BCM_ERR_OK;
+        bcmos_errno dev_res_fail = BCM_ERR_COMM_FAIL;
+
+        virtual void SetUp() {
+            bcmos_errno bal_cfg_get_stub_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.topology.topology_maps.len = num_of_pon_ports;
+        }
+
+        virtual void TearDown() {
+        }
+};
+
+// Test 1 - If querying the OLT fails, the method must return error
+TEST_F(TestProbeDevCapabilities, ProbeDev_OltQueryFailed) {
+
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__olt_topology_stub, bcmolt_cfg_get__olt_topology_stub(_,_))
+                     .WillOnce(Return(olt_res_fail));
+
+    Status query_status = ProbeDeviceCapabilities_();
+    ASSERT_TRUE( query_status.error_message() != Status::OK.error_message() );
+}
+
+// Test 2 - If all devices are queried successfully, the method must return Status::OK
+TEST_F(TestProbeDevCapabilities, ProbeDev_OltQuerySucceeded_DevQueriesSucceeded) {
+
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__olt_topology_stub, bcmolt_cfg_get__olt_topology_stub(_,_))
+                     .WillOnce(Return(olt_res_success));
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+        .WillRepeatedly(Return(dev_res_success));
+
+    Status query_status = ProbeDeviceCapabilities_();
+
+    ASSERT_TRUE( query_status.error_message() == Status::OK.error_message() );
+}
+
+// Test 3 - After successfully probing the OLT, even if probing all devices failed, the method must return error
+TEST_F(TestProbeDevCapabilities, ProbedDev_OltQuerySucceeded_AllDevQueriesFailed) {
+
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__olt_topology_stub, bcmolt_cfg_get__olt_topology_stub(_,_))
+                     .WillOnce(Return(olt_res_success));
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+        .WillRepeatedly(Return(dev_res_fail));
+
+    Status query_status = ProbeDeviceCapabilities_();
+
+    ASSERT_TRUE( query_status.error_message() != Status::OK.error_message() );
+}
+
+// Test 4 - After successfully probing the OLT, if probing some devices fail, the method returns success
+TEST_F(TestProbeDevCapabilities, ProbedDev_OltQuerySucceeded_SomeDevQueriesFailed) {
+
+    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__olt_topology_stub, bcmolt_cfg_get__olt_topology_stub(_,_))
+                     .WillOnce(Return(olt_res_success));
+
+    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
+        .WillOnce(Return(olt_res_success))
+        .WillRepeatedly(Return(dev_res_fail));
+
+    Status query_status = ProbeDeviceCapabilities_();
+
+    ASSERT_TRUE( query_status.error_message() == Status::OK.error_message() );
+}
+
diff --git a/agent/test/src/test_disable_reenable_olt.cc b/agent/test/src/test_disable_reenable_olt.cc
deleted file mode 100644
index 01fa851..0000000
--- a/agent/test/src/test_disable_reenable_olt.cc
+++ /dev/null
@@ -1,134 +0,0 @@
-/*
- * Copyright 2018-present Open Networking Foundation
-
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
-
- * http://www.apache.org/licenses/LICENSE-2.0
-
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "gtest/gtest.h"
-#include "bal_mocker.h"
-#include "core.h"
-using namespace testing;
-using namespace std;
-int num_of_pon_ports = 16;
-
-// This is used to set custom bcmolt_cfg value to bcmolt_cfg pointer coming in
-// .bcmolt_cfg_get__olt_topology_stub
-ACTION_P(SetArg1ToBcmOltCfg, value) {*static_cast<bcmolt_olt_cfg*>(arg1) = value;};
-
-
-// Create a mock function for bcmolt_cfg_get__olt_topology_stub C++ function
-MOCK_GLOBAL_FUNC2(bcmolt_cfg_get__olt_topology_stub, bcmos_errno(bcmolt_oltid, void*));
-
-class TestOltDisableReenable : public Test {
- protected:
-  virtual void SetUp() {
-    NiceMock<BalMocker> balMock;
-    bcmos_errno bal_cfg_get_stub_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.topology.topology_maps.len = num_of_pon_ports;
-    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__olt_topology_stub, bcmolt_cfg_get__olt_topology_stub(_, _))
-                     .WillOnce(DoAll(SetArg1ToBcmOltCfg(olt_cfg), Return(bal_cfg_get_stub_res)));
-
-    ProbeDeviceCapabilities_();
-
-  }
-
-  virtual void TearDown() {
-    // Code here will be called immediately after each test
-    // (right before the destructor).
-  }
-};
-
-
-// Test Fixture for OltDisable
-
-// Test 1: OltDisableSuccess  case
-TEST_F(TestOltDisableReenable, OltDisableSuccess){
-    // 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 bal_cfg_get_res = BCM_ERR_OK;
-
-    Status olt_disable_res;
-
-    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
-                         .Times(num_of_pon_ports)
-                         .WillRepeatedly(Return(bal_cfg_get_res));
-
-    olt_disable_res = Disable_();
-    ASSERT_TRUE( olt_disable_res.error_message() == Status::OK.error_message() );
-
-}
-
-// Test 2: OltDisableAllPonFailed  case
-TEST_F(TestOltDisableReenable, OltDisableAllPonFailed){
-    // 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 bal_cfg_get_res = BCM_ERR_INTERNAL;
-
-    Status olt_disable_res;
-
-    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
-                         .Times(num_of_pon_ports)
-                         .WillRepeatedly(Return(bal_cfg_get_res));
-
-    olt_disable_res = Disable_();
-    ASSERT_TRUE( olt_disable_res.error_code() == grpc::StatusCode::INTERNAL);
-}
-
-
-// Test Fixture for OltReenable
-
-// Test 1: OltReenableSuccess  case
-TEST_F(TestOltDisableReenable, OltReenableSuccess){
-    // 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 bal_cfg_get_res = BCM_ERR_OK;
-
-    Status olt_reenable_res;
-
-    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
-                         .Times(num_of_pon_ports*2)
-                         .WillRepeatedly(Return(bal_cfg_get_res));
-
-    olt_reenable_res = Reenable_();
-    ASSERT_TRUE( olt_reenable_res.error_message() == Status::OK.error_message() );
-
-}
-
-// Test 2: OltReenableAllPonFailed  case
-TEST_F(TestOltDisableReenable, OltReenableAllPonFailed){
-    // 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 bal_cfg_get_res = BCM_ERR_OK;
-    bcmos_errno olt_oper_res = BCM_ERR_INTERNAL;
-
-    Status olt_reenable_res;
-
-    EXPECT_CALL(balMock, bcmolt_cfg_get(_, _))
-                         .Times(num_of_pon_ports)
-                         .WillRepeatedly(Return(bal_cfg_get_res));
-    EXPECT_CALL(balMock,bcmolt_oper_submit(_, _))
-                         .Times(num_of_pon_ports)
-                         .WillRepeatedly(Return(olt_oper_res));
-    olt_reenable_res = Reenable_();
-    ASSERT_TRUE( olt_reenable_res.error_code() == grpc::StatusCode::INTERNAL);
-}
-
diff --git a/agent/test/src/test_enable_olt.cc b/agent/test/src/test_enable_olt.cc
deleted file mode 100644
index c4e60e7..0000000
--- a/agent/test/src/test_enable_olt.cc
+++ /dev/null
@@ -1,193 +0,0 @@
-/*
- * Copyright 2018-present Open Networking Foundation
-
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
-
- * http://www.apache.org/licenses/LICENSE-2.0
-
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-#include "gtest/gtest.h"
-#include "bal_mocker.h"
-#include "core.h"
-
-using namespace testing;
-
-class TestOltEnable : public Test {
- protected:
-  virtual void SetUp() {
-  }
-
-  virtual void TearDown() {
-    // Code here will be called immediately after each test
-    // (right before the destructor).
-  }
-};
-
-// This is used to set custom bcmolt_cfg value to bcmolt_cfg pointer coming in
-// bcmolt_cfg_get__bal_state_stub.
-ACTION_P(SetArg1ToBcmOltCfg, value) { *static_cast<bcmolt_olt_cfg*>(arg1) = value; };
-
-
-// Create a mock function for bcmolt_cfg_get__bal_state_stub C++ function
-MOCK_GLOBAL_FUNC2(bcmolt_cfg_get__bal_state_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.
-    // 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_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;
-
-    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
-    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__bal_state_stub, bcmolt_cfg_get__bal_state_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 2: OltEnableFail_host_init_fail
-TEST_F(TestOltEnable, OltEnableFail_host_init_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_INTERNAL;
-
-    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));
-
-    olt_enable_res = Enable_(1, NULL);
-    ASSERT_TRUE( olt_enable_res.error_message() != Status::OK.error_message() );
-}
-
-// Test 3: OltEnableSuccess_PON_Device_Connected
-TEST_F(TestOltEnable, OltEnableSuccess_PON_Device_Connected) {
-
-    // 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_OK;
-    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;
-
-    // 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__bal_state_stub, bcmolt_cfg_get__bal_state_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__bal_state_stub, bcmolt_cfg_get__bal_state_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.
-    // 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_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();
-
-    ON_CALL(balMock, bcmolt_host_init(_)).WillByDefault(Return(host_init_res));
-    EXPECT_GLOBAL_CALL(bcmolt_cfg_get__bal_state_stub, bcmolt_cfg_get__bal_state_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(_, _))
-                         .WillOnce(Return(olt_oper_res_fail))
-                         .WillRepeatedly(Return(olt_oper_res_success));
-    olt_enable_res = Enable_(1, NULL);
-
-    ASSERT_TRUE( olt_enable_res.error_message() == Status::OK.error_message() );
-}