test_batch_size

Signed-off-by: Kevin Su <[email protected]>
flyteorg · Apr 3, 2023 · ea7e1aa · ea7e1aa
1 parent 0977021
commit ea7e1aa
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/flytekit/unit/core/test_flyte_pickle.py b/tests/flytekit/unit/core/test_flyte_pickle.py
@@ -14,7 +14,7 @@
 from flytekit.models.literals import BlobMetadata
 from flytekit.models.types import LiteralType
 from flytekit.tools.translator import get_serializable
-from flytekit.types.pickle.pickle import FlytePickle, FlytePickleTransformer
+from flytekit.types.pickle.pickle import BatchSize, FlytePickle, FlytePickleTransformer
 
 default_img = Image(name="default", fqn="test", tag="tag")
 serialization_settings = flytekit.configuration.SerializationSettings(
@@ -55,6 +55,11 @@ def test_get_literal_type():
     )
 
 
+def test_batch_size():
+    bs = BatchSize(5)
+    assert bs.val == 5
+
+
 def test_nested():
     class Foo(object):
         def __init__(self, number: int):