x

Author: ddl-rliu

tests/flytekit/unit/core/test_flyte_file.py

@@ -17,7 +17,7 @@
 from flytekit.core.hash import HashMethod
 from flytekit.core.launch_plan import LaunchPlan
 from flytekit.core.task import task
-from flytekit.core.type_engine import TypeEngine
+from flytekit.core.type_engine import FileDownloadConfig, TypeEngine
 from flytekit.core.workflow import workflow
 from flytekit.models.core.types import BlobType
 from flytekit.models.literals import LiteralMap, Blob, BlobMetadata
@@ -764,6 +764,18 @@ def test_headers():
     assert len(FlyteFilePathTransformer.get_additional_headers(".gz")) == 1
 
 
+def test_transform_flytefile_with_file_download_config():
+    csv_file_no_config = FlyteFile["csv"]
+    lt = FlyteFilePathTransformer().get_literal_type(csv_file_no_config)
+    assert lt.blob.file_extension == ""
+    assert lt.blob.enable_legacy_filename == False
+
+    legacy_file = Annotated[FlyteFile["csv"], FileDownloadConfig(file_extension="csv", enable_legacy_filename=True)]
+    lt = FlyteFilePathTransformer().get_literal_type(legacy_file)
+    assert lt.blob.file_extension == "csv"
+    assert lt.blob.enable_legacy_filename == True
+
+
 def test_new_remote_file():
     nf = FlyteFile.new_remote_file(name="foo.txt")
     assert isinstance(nf, FlyteFile)

tests/flytekit/unit/models/core/test_types.py

@@ -15,11 +15,33 @@ def test_blob_type():
     )
     assert o.format == "csv"
     assert o.dimensionality == _types.BlobType.BlobDimensionality.SINGLE
+    assert o.file_extension == ""
+    assert o.enable_legacy_filename == False
 
     o2 = _types.BlobType.from_flyte_idl(o.to_flyte_idl())
     assert o == o2
     assert o2.format == "csv"
     assert o2.dimensionality == _types.BlobType.BlobDimensionality.SINGLE
+    assert o2.file_extension == ""
+    assert o2.enable_legacy_filename == False
+
+    o = _types.BlobType(
+        format="csv",
+        dimensionality=_types.BlobType.BlobDimensionality.SINGLE,
+        file_extension="csv",
+        enable_legacy_filename=True,
+    )
+    assert o.format == "csv"
+    assert o.dimensionality == _types.BlobType.BlobDimensionality.SINGLE
+    assert o.file_extension == "csv"
+    assert o.enable_legacy_filename == True
+
+    o2 = _types.BlobType.from_flyte_idl(o.to_flyte_idl())
+    assert o == o2
+    assert o2.format == "csv"
+    assert o2.dimensionality == _types.BlobType.BlobDimensionality.SINGLE
+    assert o2.file_extension == "csv"
+    assert o2.enable_legacy_filename == True
 
 
 def test_enum_type():