def test_several_sessionized_queues(self):
        senders = []
        receivers = []
        cookies = []
        texts = []

        for i in range(10):
            cookie = "c-" + str(i)
            cookies.append(cookie)

            sender = Sender(self.conn, "/test.sessionized/several")
            sender.start()
            senders.append(sender)

            receiver = Receiver(self.conn,
                                "/test.sessionized/several",
                                cookie=cookie)
            receiver.start()
            receivers.append(receiver)

            text = "text" + str(i)
            texts.append(text)

        arr = list(range(10))[::-1]
        random.shuffle(arr)

        # Send requests in random order
        for pos in arr:
            senders[pos].send(texts[pos], headers={"COOKIE": cookies[pos]})

        for i in range(10):
            assert texts[i] == receivers[i].receive().task
예제 #2
0
    def test_register_queue(self):
        conn = WeaveConnection.local()
        conn.connect()
        client = RPCClient(conn, self.appmgr_rpc_info, TEST_APP_TOKEN)
        client.start()

        res = client["register_queue"]("test_queue/",
                                       "fifo", {
                                           "type": "string"
                                       }, [MESSAGING_SERVER_URL], [TEST_URL],
                                       _block=True)
        assert res == "/channels/{}/test_queue".format(TEST_URL)

        sender_no_auth = Sender(conn, res, auth=TEST_APP_TOKEN)
        sender_no_auth.start()
        with pytest.raises(Unauthorized):
            sender_no_auth.send("test")

        sender_auth = Sender(conn, res, auth=MESSAGING_APP_TOKEN)
        sender_auth.send("test")

        receiver_no_auth = Receiver(conn, res, auth=MESSAGING_APP_TOKEN)
        with pytest.raises(Unauthorized):
            receiver_no_auth.receive()

        receiver_auth = Receiver(conn, res, auth=TEST_APP_TOKEN)
        assert "test" == receiver_auth.receive().task

        client.stop()
    def test_multiple_push_pop(self):
        obj = {"foo": "bar"}

        s = Sender(self.conn, "/a.b.c")
        r = Receiver(self.conn, "/a.b.c")

        s.start()
        for _ in range(10):
            s.send(obj)

        expected_message_count = 10

        def on_message(msg, headers):
            assert msg == obj
            nonlocal expected_message_count
            if expected_message_count == 1:
                r.stop()
            expected_message_count -= 1

        r.start()
        r.on_message = on_message
        thread = Thread(target=r.run)
        thread.start()

        thread.join()
    def test_simple_sessionized_push_pop(self):
        sender1 = Sender(self.conn, "/test.sessionized2")
        sender1.start()
        sender1.send("test", headers={"COOKIE": "xyz"})

        sender2 = Sender(self.conn, "/test.sessionized2")
        sender2.start()
        sender2.send("diff", headers={"COOKIE": "diff"})

        msgs1 = []
        sem1 = Semaphore(0)
        receiver1 = Receiver(self.conn, "/test.sessionized2", cookie="xyz")
        receiver1.on_message = make_receiver(2, msgs1, sem1, receiver1)
        receiver1.start()
        Thread(target=receiver1.run).start()

        msgs2 = []
        sem2 = Semaphore(0)
        receiver2 = Receiver(self.conn, "/test.sessionized2", cookie="diff")
        receiver2.on_message = make_receiver(2, msgs2, sem2, receiver2)
        receiver2.start()
        Thread(target=receiver2.run).start()

        assert sem1.acquire(timeout=10)
        assert sem2.acquire(timeout=10)
        assert msgs1[0] == "test"
        assert msgs2[0] == "diff"

        # Test retrieving items for the second time.
        sender1.send("test2", headers={"COOKIE": "xyz"})
        assert sem1.acquire(timeout=10)
        assert msgs1[1] == "test2"

        assert not sem2.acquire(timeout=5)
    def test_push_sessionized_without_key(self):
        sender = Sender(self.conn, "/test.sessionized")
        sender.start()
        with pytest.raises(ProtocolError):
            sender.send("test")

        receiver = Receiver(self.conn, "/test.sessionized")
        receiver.start()

        with pytest.raises(ProtocolError):
            receiver.receive()
    def test_multicast_with_synonym(self):
        msgs = []
        sem = Semaphore(0)
        receiver = Receiver(self.conn, "/synonyms/multi")
        receiver.on_message = make_receiver(1, msgs, sem, receiver)
        receiver.start()
        Thread(target=receiver.run).start()

        sender = Sender(self.conn, "/synonyms/multi")
        sender.start()
        sender.send("test")

        assert sem.acquire(timeout=10)
        assert msgs[-1] == "test"
        assert not sem.acquire(timeout=2)
    def test_simple_push_pop(self):
        msgs = []

        s = Sender(self.conn, "/a.b.c")
        r = Receiver(self.conn, "/a.b.c")
        s.start()
        r.start()
        r.on_message = lambda msg, hdrs: msgs.append(msg) or r.stop()

        s.send({"foo": "bar"})
        thread = Thread(target=r.run)
        thread.start()

        thread.join()

        assert msgs == [{"foo": "bar"}]
    def test_fifo_push_pop(self):
        msgs1 = []
        sem1 = Semaphore(0)
        receiver1 = Receiver(self.conn, "/test.fifo/simple")
        receiver1.on_message = make_receiver(2, msgs1, sem1, receiver1)
        receiver1.start()
        Thread(target=receiver1.run).start()

        msgs2 = []
        sem2 = Semaphore(0)
        receiver2 = Receiver(self.conn, "/test.fifo/simple")
        receiver2.on_message = make_receiver(2, msgs2, sem2, receiver2)
        receiver2.start()
        Thread(target=receiver2.run).start()

        sender1 = Sender(self.conn, "/test.fifo/simple")
        sender1.start()

        sender1.send("test")

        assert sem1.acquire(timeout=10)
        assert msgs1[-1] == "test"
        assert not sem2.acquire(timeout=2)

        sender1.send("test2")

        assert sem2.acquire(timeout=10)
        assert msgs2[-1] == "test2"
        assert not sem1.acquire(timeout=2)

        sender1.send("test3")

        assert sem1.acquire(timeout=10)
        assert msgs1[-1] == "test3"
        assert not sem2.acquire(timeout=2)

        sender1.send("test4")

        assert sem2.acquire(timeout=10)
        assert msgs2[-1] == "test4"
        assert not sem1.acquire(timeout=2)
    def test_multicast(self):
        msgs = [[] for _ in range(5)]
        sems = [Semaphore(0) for _ in range(5)]
        receivers = [Receiver(self.conn, '/multicast/1') for _ in range(5)]
        for receiver, sem, msg in zip(receivers, sems, msgs):
            receiver.on_message = make_receiver(1, msg, sem, receiver)
            receiver.start()
            Thread(target=receiver.run).start()

        sender = Sender(self.conn, '/multicast/1')
        sender.start()
        sender.send("test")

        for msg, sem in zip(msgs, sems):
            assert sem.acquire(timeout=10)
            assert msg[-1] == "test"
            assert not sem.acquire(timeout=2)

        sender.close()
        for receiver in receivers:
            receiver.stop()
예제 #10
0
    def test_queue_waits_removed_after_client_disconnects(self, queue_name):
        conn1 = WeaveConnection.local()
        conn2 = WeaveConnection.local()
        conn3 = WeaveConnection.local()
        conn1.connect()
        conn2.connect()
        conn3.connect()

        msgs1 = []
        sem1 = Semaphore(0)
        receiver1 = Receiver(conn1, queue_name, cookie="a")
        receiver1.on_message = make_receiver(1, msgs1, sem1, receiver1)
        receiver1.start()
        Thread(target=receiver1.run).start()

        msgs2 = []
        sem2 = Semaphore(0)
        receiver2 = Receiver(conn2, queue_name, cookie="b")
        receiver2.on_message = make_receiver(1, msgs2, sem2, receiver2)
        receiver2.start()
        Thread(target=receiver2.run).start()

        conn1.close()

        import time
        time.sleep(1)

        sender1 = Sender(conn3, queue_name)
        sender1.start()

        sender1.send("test", headers={"COOKIE": "b"})

        assert sem2.acquire(timeout=5)
        assert msgs2[-1] == "test"

        conn2.close()
        conn3.close()
예제 #11
0
파일: rpc.py 프로젝트: HomeWeave/WeaveLib
class RPCServer(RPC):
    MAX_RPC_WORKERS = 5

    def __init__(self,
                 name,
                 description,
                 apis,
                 service,
                 allowed_requestors=None):
        if not isinstance(service, MessagingEnabled):
            raise BadArguments("Service is not messaging enabled.")

        super(RPCServer, self).__init__(name, description, apis)
        self.service = service
        self.executor = ThreadPoolExecutor(self.MAX_RPC_WORKERS)
        self.sender = None
        self.receiver = None
        self.receiver_thread = None
        self.cookie = None
        self.appmgr_client = None
        self.allowed_requestors = allowed_requestors or []

    def register_rpc(self):
        apis = {name: api.info for name, api in self.apis.items()}
        # TODO: This means one extra RC for every registration. Clean up.
        result = self.appmgr_client["register_rpc"](self.name,
                                                    self.description,
                                                    apis,
                                                    self.allowed_requestors,
                                                    _block=True)
        return result

    def update_rpc(self, callback=None):
        apis = {name: api.info for name, api in self.apis.items()}
        return self.appmgr_client["update_rpc"](self.name,
                                                apis,
                                                _block=(not callback),
                                                _callback=callback)

    def start(self):
        conn = self.service.get_connection()
        auth_token = self.service.get_auth_token()
        self.appmgr_client = self.get_appmgr_client()
        self.appmgr_client.start()

        rpc_info = self.register_rpc()
        self.sender = Sender(conn, rpc_info["response_queue"], auth=auth_token)
        self.receiver = RPCReceiver(conn,
                                    self,
                                    rpc_info["request_queue"],
                                    auth=auth_token)

        self.sender.start()
        self.receiver.start()

        self.receiver_thread = Thread(target=self.receiver.run)
        self.receiver_thread.start()

    def get_appmgr_client(self):
        # This is so that RootRPCServer in WeaveServer need not create an
        # RPCClient to itself.
        rpc_info = find_rpc(self.service, MESSAGING_PLUGIN_URL, "app_manager")
        return RPCClient(self.service.get_connection(), rpc_info,
                         self.service.get_auth_token())

    def stop(self):
        self.appmgr_client["unregister_rpc"](self.name, _block=True)

        self.appmgr_client.stop()

        # TODO: Delete the queue, too.
        self.receiver.stop()
        self.receiver_thread.join()

        self.executor.shutdown()

    def on_rpc_message(self, rpc_obj, headers):
        def make_done_callback(request_id, cmd, cookie):
            def callback(future):
                try:
                    self.sender.send(
                        {
                            "id": request_id,
                            "command": cmd,
                            "result": future.result()
                        },
                        headers={"COOKIE": cookie})
                except WeaveException as e:
                    logger.warning("WeaveException was raised by API: %s", e)
                    self.sender.send(
                        {
                            "id": request_id,
                            "command": cmd,
                            "error_name": e.err_msg(),
                            "error": e.extra
                        },
                        headers={"COOKIE": cookie})
                except Exception as e:
                    logger.exception("Internal API raised exception." + str(e))
                    self.sender.send(
                        {
                            "id": request_id,
                            "command": cmd,
                            "error": "Internal API Error."
                        },
                        headers={"COOKIE": cookie})

            return callback

        def execute_api_internal(rpc_obj, headers, api, *args, **kwargs):
            # Keep func name in sync one in get_rpc_caller(..)
            return api(*args, **kwargs)

        obj = rpc_obj["invocation"]
        cookie = rpc_obj["response_cookie"]
        request_id = obj["id"]
        cmd = obj["command"]
        try:
            api = self[cmd]
        except KeyError:
            self.sender.send({
                "id": request_id,
                "result": False,
                "error": "API not found."
            })
            return

        args = obj.get("args", [])
        kwargs = obj.get("kwargs", {})
        future = self.executor.submit(execute_api_internal, rpc_obj, headers,
                                      api, *args, **kwargs)
        future.add_done_callback(make_done_callback(request_id, cmd, cookie))

    @property
    def info_message(self):
        return {
            "name": self.name,
            "description": self.description,
            "apis": {name: api.info
                     for name, api in self.apis.items()},
            "request_queue": self.receiver.channel,
            "response_queue": self.sender.channel
        }
예제 #12
0
파일: rpc.py 프로젝트: HomeWeave/WeaveLib
class RPCClient(RPC):
    def __init__(self, conn, rpc_info, token=None):
        self.token = token
        name = rpc_info["name"]
        description = rpc_info["description"]
        apis = [self.get_api_call(x) for x in rpc_info["apis"].values()]
        super(RPCClient, self).__init__(name, description, apis)

        self.client_cookie = "rpc-client-cookie-" + str(uuid4())
        self.sender = Sender(conn, rpc_info["request_queue"], auth=self.token)
        self.receiver = RPCReceiver(conn,
                                    self,
                                    rpc_info["response_queue"],
                                    cookie=self.client_cookie)
        self.receiver_thread = Thread(target=self.receiver.run)

        self.callbacks = {}
        self.callbacks_lock = RLock()

    def start(self):
        self.sender.start()
        self.receiver.start()
        self.receiver_thread.start()

    def stop(self):
        self.sender.close()
        self.receiver.stop()
        self.receiver_thread.join()

    def get_api_call(self, obj):
        def make_blocking_callback(event, response_arr):
            def callback(obj):
                response_arr.append(obj)
                event.set()

            return callback

        def on_invoke(obj, block, callback):
            msg_id = obj["id"]

            if block:
                res_arr = []
                event = Event()
                callback = make_blocking_callback(event, res_arr)

            if callback:
                with self.callbacks_lock:
                    self.callbacks[msg_id] = callback

            self.sender.send(
                {
                    "invocation": obj,
                    "response_cookie": self.client_cookie
                },
                headers={"AUTH": self.token})
            if not block:
                return

            event.wait()
            return extract_rpc_payload(res_arr[0])

        return ClientAPI.from_info(obj, on_invoke)

    def on_rpc_message(self, msg, headers):
        with self.callbacks_lock:
            callback = self.callbacks.pop(msg["id"])

        if not callback:
            return

        callback(msg)
예제 #13
0
 def test_push_with_bad_schema(self):
     s = Sender(self.conn, "/a.b.c")
     s.start()
     with pytest.raises(SchemaValidationFailed):
         s.send({"foo": [1, 2]})
예제 #14
0
 def test_push_to_unknown_queue(self):
     s = Sender(self.conn, "unknown.queue")
     s.start()
     with pytest.raises(ObjectNotFound):
         s.send({"a": "b"})
예제 #15
0
 def test_push_without_task(self):
     s = Sender(self.conn, "/a.b.c")
     s.start()
     with pytest.raises(ProtocolError):
         s.send(None)