<commit_before>from __future__ import with_statement

from contextlib import contextmanager
import gevent
import gevent_zeromq
from gevent_selfpipe import selfpiped

from bson import BSON
from nose.tools import assert_equal
import zmq

from zrpc.concurrency import GeventCallback
from zrpc.server import Registry
from zrpc.multiserver import MultiServer


REGISTRY = Registry()

@REGISTRY.method
def add(x, y):
    return x + y


@contextmanager
def multiserver_and_client(address, registry, n_workers, context=None):
    context = context or zmq.Context.instance()

    try:
        cb = GeventCallback()
        ms = MultiServer(address, registry, context=context)
        ms_gl = cb.spawn(ms.run, args=(n_workers,),
                         kwargs={'callback': cb,
                                 'device': selfpiped(zmq.device)})

        cb.wait()
        client = context.socket(zmq.REQ)
        client.connect(address)

        try:
            yield client
        finally:
            client.close()
    finally:
        selfpiped(context.term)()


def test_server_responds_correctly():
    context = gevent_zeromq.zmq.Context()
    with multiserver_and_client('inproc://zrpc', REGISTRY, 4, context=context) as client:
        client.send(BSON.encode({
            "id": "abc",
            "method": "add",
            "params": [3, 4]}))
        assert_equal(BSON(client.recv()).decode(),
                     {"id": "abc", "result": 7, "error": None})<commit_msg>Pass lb_context into MultiServer in gevent multiserver test.<commit_after>from __future__ import with_statement

from contextlib import contextmanager
import gevent
import gevent_zeromq
from gevent_selfpipe import selfpiped

from bson import BSON
from nose.tools import assert_equal
import zmq

from zrpc.concurrency import GeventCallback
from zrpc.server import Registry
from zrpc.multiserver import MultiServer


REGISTRY = Registry()

@REGISTRY.method
def add(x, y):
    return x + y


@contextmanager
def multiserver_and_client(address, registry, n_workers, context=None,
                           lb_context=None):
    context = context or zmq.Context.instance()

    try:
        cb = GeventCallback()
        ms = MultiServer(address, registry, context=context,
                         lb_context=lb_context)
        ms_gl = cb.spawn(ms.run, args=(n_workers,),
                         kwargs={'callback': cb,
                                 'device': selfpiped(zmq.device)})

        cb.wait()
        client = context.socket(zmq.REQ)
        client.connect(address)

        try:
            yield client
        finally:
            client.close()
    finally:
        selfpiped(context.term)()


def test_server_responds_correctly():
    context = gevent_zeromq.zmq.Context()
    blocking_context = super(gevent_zeromq.zmq.Context, context)
    with multiserver_and_client('inproc://zrpc', REGISTRY, 4, context=context,
                                lb_context=blocking_context) as client:
        client.send(BSON.encode({
            "id": "abc",
            "method": "add",
            "params": [3, 4]}))
        assert_equal(BSON(client.recv()).decode(),
                     {"id": "abc", "result": 7, "error": None})
{'fraction_non_alphanumeric': 0.10233281493001556, 'fraction_numerical': 0.002488335925349922, 'mean_word_length': 2.526315789473684, 'pattern_counts': {'":': 12, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1664814', 'n_tokens_mistral': 971, 'n_tokens_neox': 937, 'n_words': 228}