Fix bug for save_as for content type and add unittest

idmtools_core/idmtools/assets/asset.py

@@ -426,7 +426,7 @@ def save_as(self, dest: str, force: bool = False):  # noqa: F811
         if self.absolute_path is not None:
             self.download_generator_hook = partial(file_content_to_generator, self.absolute_path)
         elif self.content:
-            self.download_generator_hook = partial(content_generator, self.content)
+            self.download_generator_hook = partial(content_generator, self.bytes)
         else:
             raise ValueError("Asset has no content or absolute path")
 
@@ -438,7 +438,7 @@ def save_md5_checksum(self):
         Returns:
             None
         """
-        asset = Asset(filename=f"{self.filename}.md5", content=f"{self.filename}:md5:{self.checksum}".encode())
+        asset = Asset(filename=f"{self.filename}.md5", content=f"{self.filename}:md5:{self.checksum}")
         if asset.checksum is None:
             asset.calculate_checksum()
         asset.save_as(os.path.curdir, force=True)

idmtools_core/idmtools/utils/file.py

@@ -6,7 +6,7 @@
 import os
 from io import BytesIO
 from os import DirEntry
-from typing import Iterable, Generator, List
+from typing import Iterable, Generator, List, Union
 
 
 def scan_directory(basedir: str, recursive: bool = True, ignore_directories: List[str] = None) -> Iterable[DirEntry]:
@@ -49,7 +49,7 @@ def file_content_to_generator(absolute_path, chunk_size=128) -> Generator[bytear
                 yield res
 
 
-def content_generator(content, chunk_size=128) -> Generator[bytearray, None, None]:
+def content_generator(content: Union[str, bytes], chunk_size=128) -> Generator[bytearray, None, None]:
     """
     Create a generator from file contents in chunks(useful for streaming binary data and piping).
 
@@ -60,7 +60,10 @@ def content_generator(content, chunk_size=128) -> Generator[bytearray, None, Non
     Returns:
         Generator that return bytes in chunks of size chunk_size
     """
-    content_io = BytesIO(content)
+    if isinstance(content, str):
+        content_io = BytesIO(content.encode())
+    else:
+        content_io = BytesIO(content)
     while True:
         chunk = content_io.read(chunk_size)
         if chunk == b'':

idmtools_core/tests/test_assets.py

@@ -519,6 +519,15 @@ def test_file_content_to_generator_empty_file(self, mock_file):
         chunks = list(generator)
         self.assertEqual(len(chunks), 0)
 
+    def test_save_as(self):
+        # Initialize the asset object
+        asset = Asset(filename="example.txt", content="example content")
+        expected_content = f"{asset.content}"
+        asset.save_as("example.txt", force=True)
+        with open(os.path.join(os.path.curdir, "example.txt"), "r") as f:
+            assert f.read() == expected_content
+        os.remove(os.path.join(os.path.curdir, "example.txt"))
+
 
 if __name__ == '__main__':
     unittest.main()