From 343252c48dbd33c4a82dbe50d0c84d685b603a01 Mon Sep 17 00:00:00 2001
From: Stephen Shelton <steve@brewcraft.org>
Date: Thu, 2 Jul 2020 09:40:38 -0600
Subject: [PATCH] Use HiveContext and HiveRouter properly

---
 pybind/common.hpp            |  6 ++++++
 pybind/llarp/context.cpp     | 11 +++++++++++
 pybind/llarp/router.cpp      |  2 +-
 pybind/module.cpp            |  2 ++
 test/hive/test_peer_stats.py |  4 ++--
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/pybind/common.hpp b/pybind/common.hpp
index 125944016..6d1508882 100644
--- a/pybind/common.hpp
+++ b/pybind/common.hpp
@@ -74,4 +74,10 @@ namespace tooling
 
   void
   RouterEvent_Init(py::module& mod);
+
+  void
+  HiveContext_Init(py::module& mod);
+
+  void
+  HiveRouter_Init(py::module& mod);
 }  // namespace tooling
diff --git a/pybind/llarp/context.cpp b/pybind/llarp/context.cpp
index 1ed3aa2de..3298f14da 100644
--- a/pybind/llarp/context.cpp
+++ b/pybind/llarp/context.cpp
@@ -37,4 +37,15 @@ namespace llarp
             })
         .def("CallSafe", &Context::CallSafe);
   }
+
 }  // namespace llarp
+
+namespace tooling
+{
+  void
+  HiveContext_Init(py::module& mod)
+  {
+    py::class_<tooling::HiveContext>(mod, "HiveContext")
+        .def("getRouterAsHiveRouter", &tooling::HiveContext::getRouterAsHiveRouter);
+  }
+}  // namespace tooling
diff --git a/pybind/llarp/router.cpp b/pybind/llarp/router.cpp
index d6d74e40d..ad668f4b8 100644
--- a/pybind/llarp/router.cpp
+++ b/pybind/llarp/router.cpp
@@ -21,7 +21,7 @@ namespace tooling
   void
   HiveRouter_Init(py::module& mod)
   {
-    py::class_<HiveRouter>(mod, "HiveRouter")
+    py::class_<HiveRouter, llarp::AbstractRouter>(mod, "HiveRouter")
         .def("disableGossiping", &HiveRouter::disableGossiping)
         .def("enableGossiping", &HiveRouter::enableGossiping);
   }
diff --git a/pybind/module.cpp b/pybind/module.cpp
index aa57fa25a..605d83844 100644
--- a/pybind/module.cpp
+++ b/pybind/module.cpp
@@ -6,12 +6,14 @@ PYBIND11_MODULE(pyllarp, m)
   tooling::RouterHive_Init(m);
   tooling::RouterEvent_Init(m);
   llarp::AbstractRouter_Init(m);
+  tooling::HiveRouter_Init(m);
   llarp::PeerDb_Init(m);
   llarp::PeerStats_Init(m);
   llarp::RouterID_Init(m);
   llarp::RouterContact_Init(m);
   llarp::CryptoTypes_Init(m);
   llarp::Context_Init(m);
+  tooling::HiveContext_Init(m);
   llarp::Config_Init(m);
   llarp::dht::DHTTypes_Init(m);
   llarp::PathTypes_Init(m);
diff --git a/test/hive/test_peer_stats.py b/test/hive/test_peer_stats.py
index 48cbebd52..0905822e8 100644
--- a/test/hive/test_peer_stats.py
+++ b/test/hive/test_peer_stats.py
@@ -77,10 +77,10 @@ def tally_rc_received_for_peer(hive, routerId):
 
   numFound = 0
 
-  def visit(relay):
+  def visit(context):
     nonlocal numFound
 
-    peerDb = relay.peerDb()
+    peerDb = context.getRouterAsHiveRouter().peerDb()
     stats = peerDb.getCurrentPeerStats(routerId);
 
     assert(stats.routerId == routerId)