File: test_encoder.py

package info (click to toggle)
python-beanie 2.0.0-1
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 1,480 kB
  • sloc: python: 14,427; makefile: 7; sh: 6
file content (203 lines) | stat: -rw-r--r-- 6,288 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
import re
from datetime import date, datetime
from enum import Enum
from uuid import uuid4

import pytest
from bson import Binary, Regex
from pydantic import AnyUrl

from beanie.odm.utils.encoder import Encoder
from beanie.odm.utils.pydantic import IS_PYDANTIC_V2
from tests.odm.models import (
    BsonRegexDoc,
    Child,
    DictEnum,
    DocumentForEncodingTest,
    DocumentForEncodingTestDate,
    DocumentWithComplexDictKey,
    DocumentWithDecimalField,
    DocumentWithEnumKeysDict,
    DocumentWithHttpUrlField,
    DocumentWithKeepNullsFalse,
    DocumentWithStringField,
    ModelWithOptionalField,
    NativeRegexDoc,
    SampleWithMutableObjects,
)


async def test_encode_datetime():
    assert isinstance(Encoder().encode(datetime.now()), datetime)

    doc = DocumentForEncodingTest(datetime_field=datetime.now())
    await doc.insert()
    new_doc = await DocumentForEncodingTest.get(doc.id)
    assert isinstance(new_doc.datetime_field, datetime)


async def test_encode_date():
    assert isinstance(Encoder().encode(datetime.now()), datetime)

    doc = DocumentForEncodingTestDate()
    await doc.insert()
    new_doc = await DocumentForEncodingTestDate.get(doc.id)
    assert new_doc.date_field == doc.date_field
    assert isinstance(new_doc.date_field, date)


async def test_encode_regex():
    raw_regex = r"^AA.*CC$"
    case_sensitive_regex = re.compile(raw_regex)
    case_insensitive_regex = re.compile(raw_regex, re.I)

    assert isinstance(Encoder().encode(case_sensitive_regex), Regex)
    assert isinstance(Encoder().encode(case_insensitive_regex), Regex)

    matching_doc = DocumentWithStringField(string_field="AABBCC")
    ignore_case_matching_doc = DocumentWithStringField(string_field="aabbcc")
    non_matching_doc = DocumentWithStringField(string_field="abc")

    for doc in (matching_doc, ignore_case_matching_doc, non_matching_doc):
        await doc.insert()

    assert {matching_doc.id, ignore_case_matching_doc.id} == {
        doc.id
        async for doc in DocumentWithStringField.find(
            DocumentWithStringField.string_field == case_insensitive_regex
        )
    }
    assert {matching_doc.id} == {
        doc.id
        async for doc in DocumentWithStringField.find(
            DocumentWithStringField.string_field == case_sensitive_regex
        )
    }


def test_encode_with_custom_encoder():
    assert isinstance(
        Encoder(custom_encoders={datetime: str}).encode(datetime.now()), str
    )


async def test_bytes():
    encoded_b = Encoder().encode(b"test")
    assert isinstance(encoded_b, Binary)
    assert encoded_b.subtype == 0

    doc = DocumentForEncodingTest(bytes_field=b"test")
    await doc.insert()
    new_doc = await DocumentForEncodingTest.get(doc.id)
    assert isinstance(new_doc.bytes_field, bytes)


async def test_bytes_already_binary():
    b = Binary(b"123", 3)
    encoded_b = Encoder().encode(b)
    assert isinstance(encoded_b, Binary)
    assert encoded_b.subtype == 3


async def test_mutable_objects_on_save():
    instance = SampleWithMutableObjects(
        d={"Bar": Child(child_field="Foo")}, lst=[Child(child_field="Bar")]
    )
    await instance.save()
    assert isinstance(instance.d["Bar"], Child)
    assert isinstance(instance.lst[0], Child)


async def test_decimal():
    test_amts = DocumentWithDecimalField(amt=1, other_amt=2)
    await test_amts.insert()
    obj = await DocumentWithDecimalField.get(test_amts.id)
    assert obj.amt == 1
    assert obj.other_amt == 2

    test_amts.amt = 6
    await test_amts.save_changes()

    obj = await DocumentWithDecimalField.get(test_amts.id)
    assert obj.amt == 6

    test_amts = (await DocumentWithDecimalField.find_all().to_list())[0]
    test_amts.other_amt = 7
    await test_amts.save_changes()

    obj = await DocumentWithDecimalField.get(test_amts.id)
    assert obj.other_amt == 7


def test_keep_nulls_false():
    model = ModelWithOptionalField(i=10)
    doc = DocumentWithKeepNullsFalse(m=model)

    encoder = Encoder(keep_nulls=False, to_db=True)
    encoded_doc = encoder.encode(doc)
    assert encoded_doc == {"m": {"i": 10}}


@pytest.mark.skipif(not IS_PYDANTIC_V2, reason="Test only for Pydantic v2")
def test_should_encode_pydantic_v2_url_correctly():
    url = AnyUrl("https://example.com")
    encoder = Encoder()
    encoded_url = encoder.encode(url)
    assert isinstance(encoded_url, str)
    # pydantic2 add trailing slash for naked url. see https://github.com/pydantic/pydantic/issues/6943
    assert encoded_url == "https://example.com/"


async def test_should_be_able_to_save_retrieve_doc_with_url():
    doc = DocumentWithHttpUrlField(url_field="https://example.com")
    assert isinstance(doc.url_field, AnyUrl)
    await doc.save()

    new_doc = await DocumentWithHttpUrlField.find_one(
        DocumentWithHttpUrlField.id == doc.id
    )

    assert isinstance(new_doc.url_field, AnyUrl)
    assert new_doc.url_field == doc.url_field


async def test_dict_with_complex_key():
    assert isinstance(Encoder().encode({uuid4(): datetime.now()}), dict)

    uuid = uuid4()
    # reset microseconds, because it looses by mongo
    dt = datetime.now().replace(microsecond=0)

    doc = DocumentWithComplexDictKey(dict_field={uuid: dt})
    await doc.insert()
    new_doc = await DocumentWithComplexDictKey.get(doc.id)

    assert isinstance(new_doc.dict_field, dict)
    assert new_doc.dict_field.get(uuid) == dt


async def test_dict_with_enum_keys():
    doc = DocumentWithEnumKeysDict(color={DictEnum.RED: "favorite"})
    await doc.save()

    assert isinstance(doc.color, dict)

    for key in doc.color:
        assert isinstance(key, Enum)
        assert key == DictEnum.RED


async def test_native_regex():
    regex = re.compile(r"^1?$|^(11+?)\1+$", (re.I | re.M | re.S) ^ re.UNICODE)
    doc = await NativeRegexDoc(regex=regex).insert()
    new_doc = await NativeRegexDoc.get(doc.id)
    assert new_doc.regex == regex
    assert new_doc.regex.pattern == r"^1?$|^(11+?)\1+$"
    assert new_doc.regex.flags == int(re.I | re.M | re.S ^ re.UNICODE)


async def test_bson_regex():
    regex = Regex(r"^1?$|^(11+?)\1+$")
    doc = await BsonRegexDoc(regex=regex).insert()
    new_doc = await BsonRegexDoc.get(doc.id)
    assert new_doc.regex == Regex(pattern=r"^1?$|^(11+?)\1+$")