<commit_before>import unittest
import io
from .transports import StdioTransport, TCPTransport
import time


def json_rpc_message(payload: str) -> bytes:
    content_length = len(payload)
    return b'Content-Length: ' + bytes(
        str(content_length), 'utf-8') + b'\r\n\r\n' + bytes(payload, 'utf-8')


class FakeProcess(object):
    def __init__(self):
        self.stdin = io.BytesIO(b'foo\nbaz\n')  # io.BufferedReader()
        self.stdout = io.BytesIO(
            json_rpc_message("hello") +
            json_rpc_message("world"))  # io.BufferedWriter()
        self.returncode = None

    def poll(self):
        return self.returncode

    def exit(self, returncode):
        self.returncode = returncode


class FakeSocket(object):
    def __init__(self, data: bytes) -> None:
        self.data = data
        self.index = 0

    def recv(self, length: int) -> bytes:
        slc = self.data[self.index:length]
        self.index = max(len(self.data), self.index + length)
        return slc

    def sendall(self, payload: bytes) -> None:
        pass


class StdioTransportTests(unittest.TestCase):
    def test_read_messages(self):

        process = FakeProcess()
        t = StdioTransport(process)  # type: ignore
        self.assertIsNotNone(t)
        received = []

        def on_receive(msg):
            received.append(msg)

        def on_close():
            pass

        t.start(on_receive, on_close)
        time.sleep(0.01)
        self.assertEqual(received, ["hello", "world"])
        t.close()


class TCPTransportTests(unittest.TestCase):
    def test_read_messages(self):
        sock = FakeSocket(
            json_rpc_message("hello") + json_rpc_message("world"))
        t = TCPTransport(sock)
        self.assertIsNotNone(t)
        received = []

        def on_receive(msg):
            received.append(msg)

        def on_close():
            pass

        t.start(on_receive, on_close)
        time.sleep(0.01)
        self.assertEqual(received, ["hello", "world"])
        t.close()<commit_msg>Add tests for writing on both transports<commit_after>import unittest
import io
from .transports import StdioTransport, TCPTransport
import time
try:
    from typing import List
    assert List
except ImportError:
    pass


def json_rpc_message(payload: str) -> bytes:
    content_length = len(payload)
    return b'Content-Length: ' + bytes(
        str(content_length), 'utf-8') + b'\r\n\r\n' + bytes(payload, 'utf-8')


class FakeProcess(object):
    def __init__(self):
        self.stdin = io.BytesIO(b'')  # io.BufferedReader()
        self.stdout = io.BytesIO(
            json_rpc_message("hello") +
            json_rpc_message("world"))  # io.BufferedWriter()
        self.returncode = None

    def poll(self):
        return self.returncode

    def exit(self, returncode):
        self.returncode = returncode


class FakeSocket(object):
    def __init__(self, received: bytes) -> None:
        self.received = received
        self.sent = []  # type: List[str]
        self.index = 0

    def recv(self, length: int) -> bytes:
        slc = self.received[self.index:length]
        if slc:
            self.index = max(len(self.received), self.index + length)
            return slc
        else:
            time.sleep(1)  # simulate blocking for the duration of the test.
            return b''

    def sendall(self, payload: str) -> None:
        self.sent.append(payload)


class StdioTransportTests(unittest.TestCase):
    def test_read_messages(self):

        process = FakeProcess()
        t = StdioTransport(process)  # type: ignore
        self.assertIsNotNone(t)
        received = []

        def on_receive(msg):
            received.append(msg)

        def on_close():
            pass

        t.start(on_receive, on_close)
        time.sleep(0.01)
        self.assertEqual(received, ["hello", "world"])
        t.close()

    def test_write_messages(self):

        process = FakeProcess()
        t = StdioTransport(process)  # type: ignore
        self.assertIsNotNone(t)
        received = []

        def on_receive(msg):
            received.append(msg)

        def on_close():
            pass

        t.start(on_receive, on_close)
        t.send("hello")
        t.send("world")
        time.sleep(0.01)
        self.assertEqual(process.stdin.getvalue(), b"helloworld")
        t.close()


class TCPTransportTests(unittest.TestCase):
    def test_read_messages(self):
        sock = FakeSocket(
            json_rpc_message("hello") + json_rpc_message("world"))
        t = TCPTransport(sock)
        self.assertIsNotNone(t)
        received = []

        def on_receive(msg):
            received.append(msg)

        def on_close():
            pass

        t.start(on_receive, on_close)
        time.sleep(0.01)
        self.assertEqual(received, ["hello", "world"])
        t.close()

    def test_write_messages(self):
        sock = FakeSocket(b'')
        t = TCPTransport(sock)
        self.assertIsNotNone(t)
        received = []

        def on_receive(msg):
            received.append(msg)

        def on_close():
            pass

        t.start(on_receive, on_close)
        # TODO: move building payload into transport instead of client.
        t.send("hello")
        t.send("world")
        time.sleep(0.1)
        self.assertEqual(sock.sent, [b"hello", b"world"])
        t.close()
{'fraction_non_alphanumeric': 0.11692650334075724, 'fraction_numerical': 0.004454342984409799, 'mean_word_length': 2.2058298631766804, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7634915', 'n_tokens_mistral': 1690, 'n_tokens_neox': 1621, 'n_words': 344}