1# Copyright 2017-2020 Palantir Technologies, Inc.
2# Copyright 2021- Python Language Server Contributors.
3
4# pylint: disable=redefined-outer-name
5
6from io import BytesIO
7import datetime
8import sys
9from unittest import mock
10import pytest
11
12from pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter
13
14
15@pytest.fixture()
16def rfile():
17    return BytesIO()
18
19
20@pytest.fixture()
21def wfile():
22    return BytesIO()
23
24
25@pytest.fixture()
26def reader(rfile):
27    return JsonRpcStreamReader(rfile)
28
29
30@pytest.fixture()
31def writer(wfile):
32    return JsonRpcStreamWriter(wfile, sort_keys=True)
33
34
35def test_reader(rfile, reader):
36    rfile.write(
37        b'Content-Length: 49\r\n'
38        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
39        b'\r\n'
40        b'{"id": "hello", "method": "method", "params": {}}'
41    )
42    rfile.seek(0)
43
44    consumer = mock.Mock()
45    reader.listen(consumer)
46
47    consumer.assert_called_once_with({
48        'id': 'hello',
49        'method': 'method',
50        'params': {}
51    })
52
53
54def test_reader_bad_message(rfile, reader):
55    rfile.write(b'Hello world')
56    rfile.seek(0)
57
58    # Ensure the listener doesn't throw
59    consumer = mock.Mock()
60    reader.listen(consumer)
61    consumer.assert_not_called()
62
63
64def test_reader_bad_json(rfile, reader):
65    rfile.write(
66        b'Content-Length: 8\r\n'
67        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
68        b'\r\n'
69        b'{hello}}'
70    )
71    rfile.seek(0)
72
73    # Ensure the listener doesn't throw
74    consumer = mock.Mock()
75    reader.listen(consumer)
76    consumer.assert_not_called()
77
78
79def test_writer(wfile, writer):
80    writer.write({
81        'id': 'hello',
82        'method': 'method',
83        'params': {}
84    })
85
86    assert wfile.getvalue() == (
87        b'Content-Length: 44\r\n'
88        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
89        b'\r\n'
90        b'{"id":"hello","method":"method","params":{}}'
91    )
92
93
94class JsonDatetime(datetime.datetime):
95    """Monkey path json datetime."""
96    def __json__(self):
97        if sys.version_info.major == 3:
98            dif = int(self.timestamp())
99        else:
100            dif = int((self - datetime.datetime(1970, 1, 1)).total_seconds())
101        return '{0}'.format(dif)
102
103
104def test_writer_bad_message(wfile, writer):
105    # A datetime isn't serializable(or poorly serializable),
106    # ensure the write method doesn't throw, but the result could be empty
107    # or the correct datetime
108    datetime.datetime = JsonDatetime
109    writer.write(datetime.datetime(
110        year=2019,
111        month=1,
112        day=1,
113        hour=1,
114        minute=1,
115        second=1,
116    ))
117
118    assert wfile.getvalue() in [
119        b'',
120        b'Content-Length: 10\r\n'
121        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
122        b'\r\n'
123        b'1546304461',
124        b'Content-Length: 10\r\n'
125        b'Content-Type: application/vscode-jsonrpc; charset=utf8\r\n'
126        b'\r\n'
127        b'1546322461'
128    ]
129