From c67119501be8d229c4f838f761c9253f1d71590c Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Mon, 18 Sep 2017 17:03:30 +0530
Subject: [PATCH] Keep allocate_id() in the enhanced router class.

---
 mitogen/core.py             |  5 -----
 mitogen/master.py           | 13 +++++++++++++
 tests/data/id_allocation.py |  8 +++++---
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index c2c5fcb1..60198973 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -868,11 +868,6 @@ class Router(object):
 
         return handle
 
-    def allocate_id(self):
-        master = Context(self, 0)
-        reply = master.send_await(Message(dst_id=0, handle=ALLOCATE_ID))
-        return reply.unpickle()
-
     def on_shutdown(self, broker):
         """Called during :py:meth:`Broker.shutdown`, informs callbacks
         registered with :py:meth:`add_handle_cb` the connection is dead."""
diff --git a/mitogen/master.py b/mitogen/master.py
index 4d17eed1..b21a8478 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -739,6 +739,7 @@ METHOD_NAMES = {
 def upgrade_router(econtext):
     if not isinstance(econtext.router, Router):  # TODO
         econtext.router.__class__ = Router  # TODO
+        econtext.router.id_allocator = ChildIdAllocator(econtext.router)
         LOG.debug('_proxy_connect(): constructing ModuleForwarder')
         ModuleForwarder(econtext.router, econtext.parent, econtext.importer)
 
@@ -785,6 +786,18 @@ class IdAllocator(object):
         )
 
 
+class ChildIdAllocator(object):
+    def __init__(self, router):
+        self.router = router
+
+    def allocate(self):
+        master = Context(self.router, 0)
+        reply = master.send_await(
+            mitogen.core.Message(dst_id=0, handle=mitogen.core.ALLOCATE_ID)
+        )
+        return reply.unpickle()
+
+
 class Router(mitogen.core.Router):
     debug = False
 
diff --git a/tests/data/id_allocation.py b/tests/data/id_allocation.py
index dcfd5230..ce83f3fc 100644
--- a/tests/data/id_allocation.py
+++ b/tests/data/id_allocation.py
@@ -1,8 +1,10 @@
 
 import mitogen.core
+import mitogen.master
 
 
-@mitogen.core.takes_router
-def allocate_an_id(router):
-    return router.allocate_id()
+@mitogen.core.takes_econtext
+def allocate_an_id(econtext):
+    mitogen.master.upgrade_router(econtext)
+    return econtext.router.allocate_id()