Allow devices to override connection closing

This commit is contained in:
matejcik
2018-12-10 16:07:43 +01:00
parent 8041ed883f
commit b30e6a8408
2 changed files with 4 additions and 5 deletions

View File

@@ -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)

View File

@@ -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