| 
 | 1 | +from datetime import datetime, timezone  | 
 | 2 | +from typing import Any  | 
 | 3 | + | 
 | 4 | +import pytest  | 
 | 5 | +from pydantic import BaseModel, ValidationError  | 
 | 6 | + | 
 | 7 | +from pydantic_extra_types.ulid import ULID  | 
 | 8 | + | 
 | 9 | +try:  | 
 | 10 | +    from ulid import ULID as _ULID  | 
 | 11 | +except ModuleNotFoundError:  # pragma: no cover  | 
 | 12 | +    raise RuntimeError(  | 
 | 13 | +        'The `ulid` module requires "python-ulid" to be installed. You can install it with "pip install python-ulid".'  | 
 | 14 | +    )  | 
 | 15 | + | 
 | 16 | + | 
 | 17 | +class Something(BaseModel):  | 
 | 18 | +    ulid: ULID  | 
 | 19 | + | 
 | 20 | + | 
 | 21 | +@pytest.mark.parametrize(  | 
 | 22 | +    'ulid, result, valid',  | 
 | 23 | +    [  | 
 | 24 | +        # Valid ULID for str format  | 
 | 25 | +        ('01BTGNYV6HRNK8K8VKZASZCFPE', '01BTGNYV6HRNK8K8VKZASZCFPE', True),  | 
 | 26 | +        ('01BTGNYV6HRNK8K8VKZASZCFPF', '01BTGNYV6HRNK8K8VKZASZCFPF', True),  | 
 | 27 | +        # Invalid ULID for str format  | 
 | 28 | +        ('01BTGNYV6HRNK8K8VKZASZCFP', None, False),  # Invalid ULID (short length)  | 
 | 29 | +        ('01BTGNYV6HRNK8K8VKZASZCFPEA', None, False),  # Invalid ULID (long length)  | 
 | 30 | +        # Valid ULID for _ULID format  | 
 | 31 | +        (_ULID.from_str('01BTGNYV6HRNK8K8VKZASZCFPE'), '01BTGNYV6HRNK8K8VKZASZCFPE', True),  | 
 | 32 | +        (_ULID.from_str('01BTGNYV6HRNK8K8VKZASZCFPF'), '01BTGNYV6HRNK8K8VKZASZCFPF', True),  | 
 | 33 | +        # Invalid _ULID for bytes format  | 
 | 34 | +        (b'\x01\xBA\x1E\xB2\x8A\x9F\xFAy\x10\xD5\xA5k\xC8', None, False),  # Invalid ULID (short length)  | 
 | 35 | +        (b'\x01\xBA\x1E\xB2\x8A\x9F\xFAy\x10\xD5\xA5k\xC8\xB6\x00', None, False),  # Invalid ULID (long length)  | 
 | 36 | +        # Valid ULID for int format  | 
 | 37 | +        (109667145845879622871206540411193812282, '2JG4FVY7N8XS4GFVHPXGJZ8S9T', True),  | 
 | 38 | +        (109667145845879622871206540411193812283, '2JG4FVY7N8XS4GFVHPXGJZ8S9V', True),  | 
 | 39 | +        (109667145845879622871206540411193812284, '2JG4FVY7N8XS4GFVHPXGJZ8S9W', True),  | 
 | 40 | +    ],  | 
 | 41 | +)  | 
 | 42 | +def test_format_for_ulid(ulid: Any, result: Any, valid: bool):  | 
 | 43 | +    if valid:  | 
 | 44 | +        assert str(Something(ulid=ulid).ulid) == result  | 
 | 45 | +    else:  | 
 | 46 | +        with pytest.raises(ValidationError, match='format'):  | 
 | 47 | +            Something(ulid=ulid)  | 
 | 48 | + | 
 | 49 | + | 
 | 50 | +def test_property_for_ulid():  | 
 | 51 | +    ulid = Something(ulid='01BTGNYV6HRNK8K8VKZASZCFPE').ulid  | 
 | 52 | +    assert ulid.hex == '015ea15f6cd1c56689a373fab3f63ece'  | 
 | 53 | +    assert ulid == '01BTGNYV6HRNK8K8VKZASZCFPE'  | 
 | 54 | +    assert ulid.datetime == datetime(2017, 9, 20, 22, 18, 59, 153000, tzinfo=timezone.utc)  | 
 | 55 | +    assert ulid.timestamp == 1505945939.153  | 
 | 56 | + | 
 | 57 | + | 
 | 58 | +def test_json_schema():  | 
 | 59 | +    assert Something.model_json_schema(mode='validation') == {  | 
 | 60 | +        'properties': {  | 
 | 61 | +            'ulid': {  | 
 | 62 | +                'anyOf': [{'type': 'integer'}, {'format': 'binary', 'type': 'string'}, {'type': 'string'}],  | 
 | 63 | +                'title': 'Ulid',  | 
 | 64 | +            }  | 
 | 65 | +        },  | 
 | 66 | +        'required': ['ulid'],  | 
 | 67 | +        'title': 'Something',  | 
 | 68 | +        'type': 'object',  | 
 | 69 | +    }  | 
 | 70 | +    assert Something.model_json_schema(mode='serialization') == {  | 
 | 71 | +        'properties': {  | 
 | 72 | +            'ulid': {  | 
 | 73 | +                'anyOf': [{'type': 'integer'}, {'format': 'binary', 'type': 'string'}, {'type': 'string'}],  | 
 | 74 | +                'title': 'Ulid',  | 
 | 75 | +            }  | 
 | 76 | +        },  | 
 | 77 | +        'required': ['ulid'],  | 
 | 78 | +        'title': 'Something',  | 
 | 79 | +        'type': 'object',  | 
 | 80 | +    }  | 
0 commit comments