From 5a08fc5421f039f5d90682bb67245cb0971c5e20 Mon Sep 17 00:00:00 2001
From: Ananth Suryanarayana <anantha@juniper.net>
Date: Fri, 31 Jul 2015 10:07:25 -0700
Subject: [PATCH] Add a new API NetworkAgentMock::RouteNextHopCount to get
 nexthops count

NetworkAgentMock::RouteNextHopCount() returns the number of routes associated
with a given prefix in a given instance. It takes necessary mutex to ensure that
the table is not modified when the route is looked up and accessed.

Change-Id: Icaf33288d851645a6240abf2fe5b95e4f1b0b45b
Partial-Bug: 1464016
---
 src/bgp/test/bgp_stress_test.cc             | 13 ++----------
 src/control-node/test/network_agent_mock.cc | 23 ++++++++++++++++++---
 src/control-node/test/network_agent_mock.h  |  5 ++++-
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/src/bgp/test/bgp_stress_test.cc b/src/bgp/test/bgp_stress_test.cc
index ab16301fde8..3508cd3f2a5 100644
--- a/src/bgp/test/bgp_stress_test.cc
+++ b/src/bgp/test/bgp_stress_test.cc
@@ -1793,20 +1793,11 @@ void BgpStressTest::VerifyXmppRouteNextHops() {
             }
             for (int rt = 0; rt < n_routes_; ++rt) {
                 Ip4Prefix prefix = GetAgentRoute(agent_id, i, rt);
-                TASK_UTIL_EXPECT_NE(
-                    static_cast<test::NetworkAgentMock::RouteEntry *>(NULL),
-                    agent->route_mgr_->Lookup(instance_name,
-                                              prefix.ToString()));
-                const test::NetworkAgentMock::RouteEntry *entry;
-                entry = agent->route_mgr_->Lookup(instance_name,
-                                                  prefix.ToString());
-                if (!entry) continue;
 
                 // We expect more next-hops, one from the agent and the rest
                 // the bgp peers.
-                TASK_UTIL_EXPECT_EQ((size_t) (1 + n_peers_),
-                  agent->route_mgr_->Lookup(instance_name,
-                    prefix.ToString())->entry.next_hops.next_hop.size());
+                TASK_UTIL_EXPECT_EQ((1 + n_peers_),
+                    agent->RouteNextHopCount(instance_name, prefix.ToString()));
             }
         }
     }
diff --git a/src/control-node/test/network_agent_mock.cc b/src/control-node/test/network_agent_mock.cc
index a81561d9d5d..cdbc7c5c474 100644
--- a/src/control-node/test/network_agent_mock.cc
+++ b/src/control-node/test/network_agent_mock.cc
@@ -1305,9 +1305,12 @@ void NetworkAgentMock::InstanceMgr<T>::Clear() {
 
 template<typename T>
 const T *NetworkAgentMock::InstanceMgr<T>::Lookup(const std::string &network,
-        const std::string &prefix) const {
+        const std::string &prefix, bool take_lock) const {
     typename InstanceMgr<T>::InstanceMap::const_iterator loc;
-    tbb::mutex::scoped_lock lock(parent_->get_mutex());
+    tbb::mutex::scoped_lock lock;
+
+    if (take_lock)
+        lock.acquire(parent_->get_mutex());
 
     loc = instance_map_.find(network);
     if (loc == instance_map_.end()) {
@@ -1335,7 +1338,9 @@ template void NetworkAgentMock::InstanceMgr<T>::Remove(const std::string &networ
 template int NetworkAgentMock::InstanceMgr<T>::Count(const std::string &network) const; \
 template int NetworkAgentMock::InstanceMgr<T>::Count() const; \
 template void NetworkAgentMock::InstanceMgr<T>::Clear(); \
-template const T *NetworkAgentMock::InstanceMgr<T>::Lookup(const std::string &network, const std::string &prefix) const;
+template const T *NetworkAgentMock::InstanceMgr<T>::Lookup( \
+    const std::string &network, const std::string &prefix, \
+    const bool take_lock) const;
 
 // RouteEntry is the same type as Inet6RouteEntry, used for both inet and inet6
 INSTANTIATE_INSTANCE_TEMPLATES(NetworkAgentMock::RouteEntry)
@@ -1352,6 +1357,18 @@ int NetworkAgentMock::RouteCount() const {
     return route_mgr_->Count();
 }
 
+// Return number of nexthops associated with a given route
+int NetworkAgentMock::RouteNextHopCount(const std::string &network,
+                                        const std::string &prefix) {
+    tbb::mutex::scoped_lock lock(get_mutex());
+
+    const RouteEntry *entry = route_mgr_->Lookup(network, prefix, false);
+    if (!entry)
+        return 0;
+
+    return entry->entry.next_hops.next_hop.size();
+}
+
 int NetworkAgentMock::Inet6RouteCount(const std::string &network) const {
     return inet6_route_mgr_->Count(network);
 }
diff --git a/src/control-node/test/network_agent_mock.h b/src/control-node/test/network_agent_mock.h
index aada5a024db..deaadad4ee0 100644
--- a/src/control-node/test/network_agent_mock.h
+++ b/src/control-node/test/network_agent_mock.h
@@ -284,7 +284,8 @@ class NetworkAgentMock {
             int Count() const;
             void Clear();
             const T *Lookup(const std::string &network,
-                    const std::string &prefix) const;
+                            const std::string &prefix,
+                            const bool take_lock = true) const;
 
         private:
             NetworkAgentMock *parent_;
@@ -331,6 +332,8 @@ class NetworkAgentMock {
 
     int RouteCount(const std::string &network) const;
     int RouteCount() const;
+    int RouteNextHopCount(const std::string &network,
+                          const std::string &prefix);
     const RouteEntry *RouteLookup(const std::string &network,
                                   const std::string &prefix) const {
         return route_mgr_->Lookup(network, prefix);