From 4afa5994e8e2c23e5446f1bf7357ff85f275d5f3 Mon Sep 17 00:00:00 2001
From: Julien Duponchelle <julien@duponchelle.info>
Date: Thu, 2 Jun 2016 13:44:12 +0200
Subject: [PATCH] Fix stacktrace when exiting with Ctrl + C

Fix #547
---
 gns3server/controller/__init__.py   | 7 +++++++
 gns3server/controller/compute.py    | 5 +++++
 gns3server/web/web_server.py        | 5 ++++-
 tests/controller/test_compute.py    | 6 ++++++
 tests/controller/test_controller.py | 9 +++++++++
 5 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/gns3server/controller/__init__.py b/gns3server/controller/__init__.py
index 0b3cd929..6c12c31d 100644
--- a/gns3server/controller/__init__.py
+++ b/gns3server/controller/__init__.py
@@ -125,10 +125,17 @@ class Controller:
         :param compute_id: Compute server identifier
         """
         compute = self.get_compute(compute_id)
+        yield from compute.close()
         del self._computes[compute_id]
         self.save()
         self.notification.emit("compute.deleted", compute.__json__())
 
+    @asyncio.coroutine
+    def close(self):
+        log.info("Close controller")
+        for compute in self._computes.values():
+            yield from compute.close()
+
     @property
     def notification(self):
         """
diff --git a/gns3server/controller/compute.py b/gns3server/controller/compute.py
index 0ed27a3e..1a60ef6d 100644
--- a/gns3server/controller/compute.py
+++ b/gns3server/controller/compute.py
@@ -93,6 +93,11 @@ class Compute:
         self._connected = False
         self._controller.notification.emit("compute.updated", self.__json__())
 
+    @asyncio.coroutine
+    def close(self):
+        self._connected = False
+        self._session.close()
+
     @property
     def version(self):
         """
diff --git a/gns3server/web/web_server.py b/gns3server/web/web_server.py
index 8c60d3cb..1f957e6c 100644
--- a/gns3server/web/web_server.py
+++ b/gns3server/web/web_server.py
@@ -91,6 +91,9 @@ class WebServer:
             yield from self._handler.finish_connections()
             self._handler = None
 
+        if Config.instance().get_section_config("Server").getboolean("controller"):
+            yield from Controller.instance().close()
+
         for module in MODULES:
             log.debug("Unloading module {}".format(module.__name__))
             m = module.instance()
@@ -182,6 +185,7 @@ class WebServer:
             # Add a periodic callback to give a chance to process signals on Windows
             # because asyncio.add_signal_handler() is not supported yet on that platform
             # otherwise the loop runs outside of signal module's ability to trap signals.
+
             def wakeup():
                 loop.call_later(0.5, wakeup)
             loop.call_later(0.5, wakeup)
@@ -189,7 +193,6 @@ class WebServer:
 
         server_config = Config.instance().get_section_config("Server")
 
-
         ssl_context = None
         if server_config.getboolean("ssl"):
             if sys.platform.startswith("win"):
diff --git a/tests/controller/test_compute.py b/tests/controller/test_compute.py
index 0c340c25..8275275f 100644
--- a/tests/controller/test_compute.py
+++ b/tests/controller/test_compute.py
@@ -205,6 +205,12 @@ def test_streamFile(project, async_run, compute):
     mock.assert_called_with("GET", "https://example.com:84/v2/compute/projects/{}/stream/test/titi".format(project.id), auth=None)
 
 
+def test_close(compute, async_run):
+    assert compute.connected is True
+    async_run(compute.close())
+    assert compute.connected is False
+
+
 def test_update(compute, controller, async_run):
     compute._controller._notification = MagicMock()
     compute.name = "Test"
diff --git a/tests/controller/test_controller.py b/tests/controller/test_controller.py
index 032d51b6..12ebc744 100644
--- a/tests/controller/test_controller.py
+++ b/tests/controller/test_controller.py
@@ -91,12 +91,14 @@ def test_deleteCompute(controller, controller_config_path, async_run):
     c = async_run(controller.add_compute(compute_id="test1"))
     assert len(controller.computes) == 1
     controller._notification = MagicMock()
+    c._connected = True
     async_run(controller.delete_compute("test1"))
     assert len(controller.computes) == 0
     controller._notification.emit.assert_called_with("compute.deleted", c.__json__())
     with open(controller_config_path) as f:
         data = json.load(f)
         assert len(data["computes"]) == 0
+    assert c.connected is False
 
 
 def test_addComputeConfigFile(controller, controller_config_path, async_run):
@@ -174,3 +176,10 @@ def test_getProject(controller, async_run):
     assert controller.get_project(uuid1) == project
     with pytest.raises(aiohttp.web.HTTPNotFound):
         assert controller.get_project("dsdssd")
+
+
+def test_close(controller, async_run):
+    c = async_run(controller.add_compute(compute_id="test1"))
+    c._connected = True
+    async_run(controller.close())
+    assert c.connected is False