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
|
from datetime import datetime
from enum import Enum
from typing import Union
import pytest
from bson import ObjectId
from monty.json import MSONable
from pydantic import BaseModel, Field
from maggma.api.utils import api_sanitize, serialization_helper
class SomeEnum(Enum):
A = 1
B = 2
C = 3
class Pet(MSONable):
def __init__(self, name, age):
self.name = name
self.age = age
class AnotherPet(MSONable):
def __init__(self, name, age):
self.name = name
self.age = age
class AnotherOwner(BaseModel):
name: str = Field(..., description="Owner name")
weight_or_pet: Union[float, AnotherPet] = Field(..., title="Owners weight or Pet")
class Owner(BaseModel):
name: str = Field(..., title="Owner's name")
age: int = Field(..., title="Owne'r Age")
weight: float = Field(..., title="Owner's weight")
last_updated: datetime = Field(..., title="Last updated date for this record")
pet: Pet = Field(..., title="Owner's Pet")
other: SomeEnum = Field(..., title="A enum?")
def test_api_sanitize():
# Ensure model validation fails
with pytest.raises(ValueError):
Owner()
# This should still fail validation
new_owner = api_sanitize(Owner, fields_to_leave=["Owner.name"])
with pytest.raises(ValueError):
new_owner()
new_owner(name="owner")
# These will fail if non-optional fields are not turned off
new_owner2 = api_sanitize(Owner)
new_owner() # api_sanitize is in-place
new_owner2()
Owner()
# This should fail type validation for pet
with pytest.raises(Exception):
Owner(pet="fido")
temp_pet_dict = Pet(name="fido", age=3).as_dict()
bad_pet_dict = dict(temp_pet_dict)
del bad_pet_dict["@module"]
del bad_pet_dict["@class"]
# This should fail because of bad data type
with pytest.raises(Exception):
Owner(pet=bad_pet_dict)
assert isinstance(Owner(pet=temp_pet_dict).pet, Pet)
api_sanitize(Owner, allow_dict_msonable=True)
# This should still fail because of bad data type
with pytest.raises(Exception):
Owner(pet=bad_pet_dict)
# This should work
assert isinstance(Owner(pet=temp_pet_dict).pet, dict)
# This should work evne though AnotherPet is inside the Union type
api_sanitize(AnotherOwner, allow_dict_msonable=True)
temp_pet_dict = AnotherPet(name="fido", age=3).as_dict()
assert isinstance(AnotherPet.validate_monty_v2(temp_pet_dict, None), dict)
def test_serialization_helper():
oid = ObjectId("60b7d47bb671aa7b01a2adf6")
assert serialization_helper(oid) == "60b7d47bb671aa7b01a2adf6"
@pytest.mark.xfail()
def test_serialization_helper_xfail():
oid = "test"
serialization_helper(oid)
|