vici: Move event (un-)registration to a helper method in Python library

Also make sure events are unregistered in case of exceptions in
streamed_request().
This commit is contained in:
Tobias Brunner
2015-08-17 11:12:17 +02:00
parent 65ac0851c0
commit cf4cefee42
3 changed files with 58 additions and 47 deletions
@@ -1,4 +1,6 @@
Copyright (c) 2015 Björn Schuberg Copyright (c) 2015 Björn Schuberg
Copyright (c) 2015 Martin Willi
Copyright (c) 2015 Tobias Brunner
Permission is hereby granted, free of charge, to any person obtaining a copy Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal of this software and associated documentation files (the "Software"), to deal
@@ -8,3 +8,6 @@ class SessionException(Exception):
class CommandException(Exception): class CommandException(Exception):
"""Command result exception.""" """Command result exception."""
class EventUnknownException(Exception):
"""Event unknown exception."""
@@ -1,7 +1,7 @@
import collections import collections
import socket import socket
from .exception import SessionException, CommandException from .exception import SessionException, CommandException, EventUnknownException
from .protocol import Transport, Packet, Message from .protocol import Transport, Packet, Message
@@ -215,6 +215,32 @@ class SessionHandler(object):
self.transport.send(packet) self.transport.send(packet)
return Packet.parse(self.transport.receive()) return Packet.parse(self.transport.receive())
def _register_unregister(self, event_type, register):
"""Register or unregister for the given event.
:param event_type: event to register
:type event_type: str
:param register: whether to register or unregister
:type register: bool
"""
if register:
packet = Packet.register_event(event_type)
else:
packet = Packet.unregister_event(event_type)
response = self._communicate(packet)
if response.response_type == Packet.EVENT_UNKNOWN:
raise EventUnknownException(
"Unknown event type '{event}'".format(event=event_type)
)
elif response.response_type != Packet.EVENT_CONFIRM:
raise SessionException(
"Unexpected response type {type}, "
"expected '{confirm}' (EVENT_CONFIRM)".format(
type=response.response_type,
confirm=Packet.EVENT_CONFIRM,
)
)
def request(self, command, message=None): def request(self, command, message=None):
"""Send request with an optional message. """Send request with an optional message.
@@ -265,57 +291,37 @@ class SessionHandler(object):
if message is not None: if message is not None:
message = Message.serialize(message) message = Message.serialize(message)
# subscribe to event stream self._register_unregister(event_stream_type, True);
packet = Packet.register_event(event_stream_type)
response = self._communicate(packet)
if response.response_type != Packet.EVENT_CONFIRM: try:
raise SessionException( packet = Packet.request(command, message)
"Unexpected response type {type}, " self.transport.send(packet)
"expected '{confirm}' (EVENT_CONFIRM)".format( exited = False
type=response.response_type, while True:
confirm=Packet.EVENT_CONFIRM, response = Packet.parse(self.transport.receive())
) if response.response_type == Packet.EVENT:
) if not exited:
try:
yield Message.deserialize(response.payload)
except GeneratorExit:
exited = True
pass
else:
break
# issue command, and read any event messages if response.response_type == Packet.CMD_RESPONSE:
packet = Packet.request(command, message) command_response = Message.deserialize(response.payload)
self.transport.send(packet)
exited = False
while True:
response = Packet.parse(self.transport.receive())
if response.response_type == Packet.EVENT:
if not exited:
try:
yield Message.deserialize(response.payload)
except GeneratorExit:
exited = True
pass
else: else:
break raise SessionException(
"Unexpected response type {type}, "
if response.response_type == Packet.CMD_RESPONSE: "expected '{response}' (CMD_RESPONSE)".format(
command_response = Message.deserialize(response.payload) type=response.response_type,
else: response=Packet.CMD_RESPONSE
raise SessionException( )
"Unexpected response type {type}, "
"expected '{response}' (CMD_RESPONSE)".format(
type=response.response_type,
response=Packet.CMD_RESPONSE
) )
)
# unsubscribe from event stream finally:
packet = Packet.unregister_event(event_stream_type) self._register_unregister(event_stream_type, False);
response = self._communicate(packet)
if response.response_type != Packet.EVENT_CONFIRM:
raise SessionException(
"Unexpected response type {type}, "
"expected '{confirm}' (EVENT_CONFIRM)".format(
type=response.response_type,
confirm=Packet.EVENT_CONFIRM,
)
)
# evaluate command result, if any # evaluate command result, if any
if "success" in command_response: if "success" in command_response: