From b30e6a840863673a9acfe77991bd0da4a8bd3a14 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 10 Dec 2018 16:07:43 +0100
Subject: [PATCH] Allow devices to override connection closing

---
 libagent/device/fake_device.py | 4 ----
 libagent/device/interface.py   | 5 ++++-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/libagent/device/fake_device.py b/libagent/device/fake_device.py
index 5ac25a2..37710d1 100644
--- a/libagent/device/fake_device.py
+++ b/libagent/device/fake_device.py
@@ -39,10 +39,6 @@ class FakeDevice(interface.Device):
         self.vk = self.sk.get_verifying_key()
         return self
 
-    def close(self):
-        """Close connection."""
-        self.conn = None
-
     def pubkey(self, identity, ecdh=False):
         """Return public key."""
         _verify_support(identity)
diff --git a/libagent/device/interface.py b/libagent/device/interface.py
index ede9e3d..3a72ac7 100644
--- a/libagent/device/interface.py
+++ b/libagent/device/interface.py
@@ -113,6 +113,9 @@ class Device:
         """Connect to device, otherwise raise NotFoundError."""
         raise NotImplementedError()
 
+    def close(self):
+        self.conn.close()
+
     def __enter__(self):
         """Allow usage as context manager."""
         self.conn = self.connect()
@@ -121,7 +124,7 @@ class Device:
     def __exit__(self, *args):
         """Close and mark as disconnected."""
         try:
-            self.conn.close()
+            self.close()
         except Exception as e:  # pylint: disable=broad-except
             log.exception('close failed: %s', e)
         self.conn = None