From ca4ad055d3635a42536d4aa4b3309b338e30a2cb Mon Sep 17 00:00:00 2001
From: PrettyWood <em.jolibois@gmail.com>
Date: Wed, 4 Nov 2020 19:12:52 +0100
Subject: [PATCH] test: add StrictUnion tests

---
 tests/test_types.py | 54 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/tests/test_types.py b/tests/test_types.py
index 7e4c7e0bd96..8c6ef02735c 100644
--- a/tests/test_types.py
+++ b/tests/test_types.py
@@ -22,6 +22,7 @@
     Sequence,
     Set,
     Tuple,
+    Union,
 )
 from uuid import UUID
 
@@ -52,6 +53,7 @@
     StrictFloat,
     StrictInt,
     StrictStr,
+    StrictUnion,
     ValidationError,
     conbytes,
     condecimal,
@@ -2540,3 +2542,55 @@ class Model(BaseModel):
         v: Deque[int]
 
     assert Model(v=deque((1, 2, 3))).json() == '{"v": [1, 2, 3]}'
+
+
+def test_strict_union_one():
+    class PikaModel(BaseModel):
+        v: Union[int] = None
+        strict_v: StrictUnion[int] = None
+
+    assert PikaModel.__fields__['strict_v'].type_.__name__ == 'StrictUnion[int]'
+
+    pika = PikaModel(v='1')
+    assert pika.v == 1
+
+    with pytest.raises(ValidationError) as exc_info:
+        PikaModel(strict_v='1')
+    assert exc_info.value.errors() == [
+        {
+            'loc': ('strict_v',),
+            'msg': 'type of value is not allowed',
+            'type': 'type_error.strict_union',
+            'ctx': {'allowed_types': 'int'},
+        }
+    ]
+
+
+def test_strict_union_multi():
+    class PikaModel(BaseModel):
+        v1: Union[int, bool, str]
+        v2: Union[bool, int, str]
+        v3: Union[str, int, bool]
+        strict_v: StrictUnion[int, bool, str]
+
+    assert PikaModel.__fields__['strict_v'].type_.__name__ == 'StrictUnion[int, bool, str]'
+
+    pika = PikaModel(v1=1, v2=1, v3=1, strict_v=1)
+    assert pika.dict() == {'v1': 1, 'v2': True, 'v3': '1', 'strict_v': 1}
+
+    pika = PikaModel(v1=True, v2=True, v3=True, strict_v=True)
+    assert pika.dict() == {'v1': 1, 'v2': True, 'v3': 'True', 'strict_v': True}
+
+    pika = PikaModel(v1='1', v2='1', v3='1', strict_v='1')
+    assert pika.dict() == {'v1': 1, 'v2': True, 'v3': '1', 'strict_v': '1'}
+
+    with pytest.raises(ValidationError) as exc_info:
+        PikaModel(v1=b'1', v2=b'1', v3=b'1', strict_v=b'1')
+    assert exc_info.value.errors() == [
+        {
+            'loc': ('strict_v',),
+            'msg': 'type of value is not allowed',
+            'type': 'type_error.strict_union',
+            'ctx': {'allowed_types': 'int, bool, str'},
+        }
+    ]