DUMMY_IAMALIVE = u"""
    <?xml version=\'1.0\' encoding=\'UTF-8\'?>
    <trn:Transport xmlns:trn="http://www.telescope-networks.org/xml/Transport/v1.1"
        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xsi:schemaLocation="http://telescope-networks.org/schema/Transport/v1.1
            http://www.telescope-networks.org/schema/Transport-v1.1.xsd"
        version="1.0" role="iamalive">
        <Origin>%s</Origin>
        <TimeStamp>2012-01-01T00:00:00Z</TimeStamp>
    </trn:Transport>
""" % (
    DUMMY_EVENT_IVOID.decode(),
) DUMMY_ACK = u"""
    <?xml version='1.0' encoding='UTF-8'?>
    <trn:Transport xmlns:trn="http://www.telescope-networks.org/xml/Transport/v1.1"
        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xsi:schemaLocation="http://telescope-networks.org/schema/Transport/v1.1
            http://www.telescope-networks.org/schema/Transport-v1.1.xsd"
        version="1.0" role="ack">
        <Origin>%s</Origin>
        <Response>%s</Response>
        <TimeStamp>2012-01-01T00:00:00Z</TimeStamp>
    </trn:Transport>
""" % (
    DUMMY_EVENT_IVOID.decode(),
    DUMMY_SERVICE_IVOID.decode(),
) DUMMY_NAK = u"""
    <?xml version='1.0' encoding='UTF-8'?>
    <trn:Transport xmlns:trn="http://www.telescope-networks.org/xml/Transport/v1.1"
        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xsi:schemaLocation="http://telescope-networks.org/schema/Transport/v1.1
            http://www.telescope-networks.org/schema/Transport-v1.1.xsd"
        version="1.0" role="nak">
        <Origin>%s</Origin>
        <Response>%s</Response>
        <TimeStamp>2012-01-01T00:00:00Z</TimeStamp>
    </trn:Transport>
""" % (
    DUMMY_EVENT_IVOID.decode(),
    DUMMY_SERVICE_IVOID.decode(),
) DUMMY_AUTHENTICATE_RESPONSE = partial(
    TransportMessage.authenticateresponse, DUMMY_SERVICE_IVOID, DUMMY_SERVICE_IVOID
) DUMMY_AUTHENTICATE_RESPONSE_LEGACY = u"""
    <?xml version='1.0' encoding='UTF-8'?>
    <trn:Transport xmlns:trn="http://www.telescope-networks.org/xml/Transport/v1.1"
        xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
        xsi:schemaLocation="http://telescope-networks.org/schema/Transport/v1.1
            http://www.telescope-networks.org/schema/Transport-v1.1.xsd"
        version="1.0" role="authenticate">
        <Origin>%s</Origin>
        <Response>%s</Response>
        <TimeStamp>2012-01-01T00:00:00Z</TimeStamp>
        <Meta>
            <filter type="xpath">%s</filter>
        </Meta>
    </trn:Transport>
""" % (
    DUMMY_SERVICE_IVOID.decode(),
    DUMMY_SERVICE_IVOID.decode(),
    "%s",
) DUMMY_SERVICE_IVOID = u"ivo://comet.broker/test".encode("UTF-8") class DummyEvent(object):
    def __init__(self, ivoid=DUMMY_EVENT_IVOID):
        self.attrib = {"ivorn": ivoid}
        self.raw_bytes = DUMMY_VOEVENT.replace(DUMMY_EVENT_IVOID, ivoid)
        self.element = etree.fromstring(self.raw_bytes) BCAST_TEST_INTERVAL = 3600 class VOEventBroadcaster(ElementSender):
    MAX_ALIVE_COUNT = 1  # Drop connection if peer misses too many iamalives
    MAX_OUTSTANDING_ACK = 10  # Drop connection if peer misses too many acks

    def __init__(self):
        self.filters = []

    def connectionMade(self):
        log.info("New subscriber at %s" % str(self.transport.getPeer()))
        self.factory.broadcasters.append(self)
        self.alive_count = 0
        self.send_xml(TransportMessage.authenticate(self.factory.local_ivo))
        self.outstanding_ack = 0

    def connectionLost(self, *args):
        log.info("Subscriber at %s disconnected" % str(self.transport.getPeer()))
        self.factory.broadcasters.remove(self)
        return ElementSender.connectionLost(self, *args)

    def sendIAmAlive(self):
        if self.alive_count >= self.MAX_ALIVE_COUNT:
            log.info("Peer appears to be dead; dropping connection")
            self.transport.loseConnection()
        elif self.outstanding_ack >= self.MAX_OUTSTANDING_ACK:
            log.info("Peer is not acknowledging events; dropping connection")
            self.transport.loseConnection()
        else:
            self.send_xml(TransportMessage.iamalive(self.factory.local_ivo))
            self.alive_count += 1

    def stringReceived(self, data):
        try:
            incoming = TransportMessage(data)
        except ParseError:
            log.warn("Unparsable message received")
            return

        if incoming.role == "iamalive":
            log.debug("IAmAlive received from %s" % str(self.transport.getPeer()))
            self.alive_count -= 1
        elif incoming.role == "ack":
            log.debug("Ack received from %s" % str(self.transport.getPeer()))
            self.outstanding_ack -= 1
        elif incoming.role == "nak":
            log.info(
                "Nak received from %s; terminating" % str(self.transport.getPeer())
            )
            self.transport.loseConnection()
        elif incoming.role == "authenticate":
            log.debug("Authentication received from %s" % str(self.transport.getPeer()))
            self.filters = []
            # Accept both "new-style" (<Param type="xpath-filter" />) and
            # old-style (<filter type="xpath" />) filters.
            for xpath in chain(
                [
                    elem.get("value")
                    for elem in incoming.element.findall(
                        'Meta/Param[@name="xpath-filter"]'
                    )
                ],
                [
                    elem.text
                    for elem in incoming.element.findall('Meta/filter[@type="xpath"]')
                ],
            ):
                log.info(
                    "Installing filter %s for %s"
                    % (xpath, str(self.transport.getPeer()))
                )
                try:
                    self.filters.append(ElementTree.XPath(xpath))
                except ElementTree.XPathSyntaxError:
                    log.info("Filter %s is not valid XPath" % (xpath,))
        else:
            log.warn(
                "Incomprehensible data received from %s (role=%s)"
                % (self.transport.getPeer(), incoming.role)
            )

    def send_event(self, event):
        # Check the event against our filters and, if one or more pass, then
        # we send the event to our subscriber.
        def check_filters(result):
            if not self.filters or any([value for success, value in result if success]):
                log.info(
                    "Event matches filter criteria: forwarding to %s"
                    % (str(self.transport.getPeer()),)
                )
                self.send_xml(event)
                self.outstanding_ack += 1
            else:
                log.info("Event rejected by filter")

        return defer.DeferredList(
            [deferToThread(xpath, event.element) for xpath in self.filters],
            consumeErrors=True,
        ).addCallback(check_filters) class VOEventBroadcasterFactory(ServerFactory):
    IAMALIVE_INTERVAL = 60  # Sent iamalive every IAMALIVE_INTERVAL seconds
    protocol = VOEventBroadcaster

    def __init__(self, local_ivo, test_interval):
        # test_interval is the time in seconds between sending test events to
        # the network. 0 to disable.
        self.local_ivo = local_ivo
        self.test_interval = test_interval
        self.broadcasters = []
        self.alive_loop = LoopingCall(self.sendIAmAlive)
        self.test_loop = LoopingCall(self.sendTestEvent)

    def startFactory(self):
        self.alive_loop.start(self.IAMALIVE_INTERVAL)
        if self.test_interval:
            self.test_loop.start(self.test_interval)
        return ServerFactory.startFactory(self)

    def stopFactory(self):
        self.alive_loop.stop()
        if self.test_loop.running:
            self.test_loop.stop()
        return ServerFactory.stopFactory(self)

    def sendIAmAlive(self):
        log.debug("Broadcasting iamalive")
        for broadcaster in self.broadcasters:
            broadcaster.sendIAmAlive()

    def sendTestEvent(self):
        log.debug("Broadcasting test event")
        test_event = VOEventMessage.broker_test(self.local_ivo)
        for broadcaster in self.broadcasters:
            broadcaster.send_event(test_event)
import lxml.etree as etree
from twisted.internet import reactor
from twisted.internet import task
from twisted.trial import unittest
from twisted.test import proto_helpers
from comet.testutils import (
    DUMMY_IAMALIVE,
    DUMMY_ACK,
    DUMMY_NAK,
    DUMMY_AUTHENTICATE_RESPONSE,
    DUMMY_AUTHENTICATE_RESPONSE_LEGACY,
    DUMMY_SERVICE_IVOID,
    DummyEvent,
)
from comet.service.broker import BCAST_TEST_INTERVAL
from comet.protocol.broadcaster import VOEventBroadcaster, VOEventBroadcasterFactory
# Comet VOEvent Broker.
# Tests for event broadcaster.






class DummyBroadcaster(object):
    def __init__(self):
        self.received_alive = False
        self.received_event = False

    def sendIAmAlive(self):
        self.received_alive = True

    def send_event(self, event):
        self.received_event = True


class VOEventBroadcasterFactoryTestCaseBase(unittest.TestCase):
    def setUp(self, test_interval=BCAST_TEST_INTERVAL):
        self.factory = VOEventBroadcasterFactory(
            DUMMY_SERVICE_IVOID.decode(), test_interval
        )
        self.factory.alive_loop.clock = task.Clock()
        self.factory.test_loop.clock = task.Clock()
        self.factory.broadcasters.append(DummyBroadcaster())
        self.connector = reactor.listenTCP(0, self.factory)

    def tearDown(self):
        self.connector.stopListening()


class VOEventBroadcasterFactoryTestCase(VOEventBroadcasterFactoryTestCaseBase):
    def test_protocol(self):
        self.assertEqual(self.factory.protocol, VOEventBroadcaster)

    def test_sendIamAlive(self):
        self.assertEqual(self.factory.alive_loop.running, True)
        self.factory.alive_loop.clock.advance(self.factory.IAMALIVE_INTERVAL)
        for broadcaster in self.factory.broadcasters:
            self.assertEqual(broadcaster.received_alive, True)

    def test_sendTestEvent(self):
        self.assertEqual(self.factory.test_loop.running, True)
        self.factory.test_loop.clock.advance(BCAST_TEST_INTERVAL)
        for broadcaster in self.factory.broadcasters:
            self.assertEqual(broadcaster.received_event, True)


class VOEventBroadcasterFactoryNoTestEventsTestCase(
    VOEventBroadcasterFactoryTestCaseBase
):
    # In this special case, we cause the factory not to broadcast test events
    # to the network.
    def setUp(self):
        VOEventBroadcasterFactoryTestCaseBase.setUp(self, 0)

    def test_DoNotSendTestEvent(self):
        self.assertEqual(self.factory.test_loop.running, False)
        self.factory.test_loop.clock.advance(BCAST_TEST_INTERVAL)
        for broadcaster in self.factory.broadcasters:
            self.assertEqual(broadcaster.received_event, False)


class VOEventBroadcasterTestCase(unittest.TestCase):
    def setUp(self):
        self.factory = VOEventBroadcasterFactory(
            DUMMY_SERVICE_IVOID.decode(), BCAST_TEST_INTERVAL
        )
        self.factory.alive_loop.clock = task.Clock()
        self.factory.test_loop.clock = task.Clock()
        self.connector = reactor.listenTCP(0, self.factory)
        self.proto = self.factory.buildProtocol(("127.0.0.1", 0))
        self.tr = proto_helpers.StringTransportWithDisconnection()
        self.proto.makeConnection(self.tr)
        self.tr.protocol = self.proto

    def tearDown(self):
        self.connector.stopListening()

    def test_register_broadcaster(self):
        # Should now be a broadcaster registered with the factory
        self.assertIn(self.proto, self.factory.broadcasters)

    def test_sent_authenticate(self):
        received_element = etree.fromstring(self.tr.value()[4:])
        self.assertEqual("authenticate", received_element.attrib["role"])
        self.assertEqual(
            DUMMY_SERVICE_IVOID.decode(), received_element.find("Origin").text
        )

    def test_sendIAmAlive(self):
        self.tr.clear()
        init_alive_count = self.proto.alive_count
        self.factory.alive_loop.clock.advance(self.factory.IAMALIVE_INTERVAL)
        received_element = etree.fromstring(self.tr.value()[4:])
        self.assertEqual("iamalive", received_element.attrib["role"])
        self.assertEqual(
            DUMMY_SERVICE_IVOID.decode(), received_element.find("Origin").text
        )
        self.assertEqual(self.proto.alive_count - init_alive_count, 1)

    def test_alive_timeout(self):
        self.assertEqual(self.tr.connected, True)
        for x in range(self.proto.MAX_ALIVE_COUNT + 1):
            self.proto.sendIAmAlive()
        self.assertEqual(self.tr.connected, False)

    def test_lack_of_ack(self):
        self.assertEqual(self.tr.connected, True)
        self.proto.outstanding_ack = self.proto.MAX_OUTSTANDING_ACK + 1
        self.proto.sendIAmAlive()
        self.assertEqual(self.tr.connected, False)

    def test_send_event(self):
        # Note lack of filter!
        self.tr.clear()
        init_outstanding_ack = self.proto.outstanding_ack

        self.proto.send_event(DummyEvent())