def test_wrap_files_in_tarball(self):
""" Test wrap files in a tarball """
# prepare
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
temp_file_one = os.path.join(test_data_dir, 'test_data_one.tmp')
temp_file_two = os.path.join(test_data_dir, 'test_data_two.tmp')
temp_tarball = os.path.join(test_data_dir, 'test_data.tmp.tar.gz')
with open(temp_file_one, 'w') as f:
f.write('temporary file one content')
with open(temp_file_two, 'w') as f:
f.write('temporary file two content')
# check
builder = ImageBuilder(gcs_base=GCS_BASE, target_image='')
builder._wrap_files_in_tarball(temp_tarball, {
'dockerfile': temp_file_one,
'main.py': temp_file_two
})
self.assertTrue(os.path.exists(temp_tarball))
with tarfile.open(temp_tarball) as temp_tarball_handle:
temp_files = temp_tarball_handle.getmembers()
self.assertTrue(len(temp_files) == 2)
for temp_file in temp_files:
self.assertTrue(temp_file.name in ['dockerfile', 'main.py'])
# clean up
os.remove(temp_file_one)
os.remove(temp_file_two)
os.remove(temp_tarball)
def test_generate_entrypoint_python2(self):
""" Test entrypoint generation for python2"""
# prepare
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
# check
builder = ImageBuilder(gcs_base=GCS_BASE, target_image='')
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func_two, python_version='python2')
golden = '''\
def sample_component_func_two(a, b):
result = 3.45
if a == 'succ':
result = float(b + 5)
return result
def wrapper_sample_component_func_two(a,b,_output_file):
output = sample_component_func_two(str(a),int(b))
import os
os.makedirs(os.path.dirname(_output_file))
with open(_output_file, "w") as data:
data.write(str(output))
import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("a", type=str)
parser.add_argument("b", type=int)
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())
if __name__ == "__main__":
wrapper_sample_component_func_two(**args)
'''
self.assertEqual(golden, generated_codes)
def test_generate_kaniko_yaml(self):
""" Test generating the kaniko job yaml """
# prepare
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
# check
builder = ImageBuilder(gcs_base=GCS_BASE, target_image='')
generated_yaml = builder._generate_kaniko_spec(namespace='default', arc_dockerfile_name='dockerfile',
gcs_path='gs://mlpipeline/kaniko_build.tar.gz', target_image='gcr.io/mlpipeline/kaniko_image:latest')
with open(os.path.join(test_data_dir, 'kaniko.basic.yaml'), 'r') as f:
golden = yaml.load(f)
self.assertEqual(golden, generated_yaml)
def test_generate_entrypoint(self):
""" Test entrypoint generation """
# prepare
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
# check
builder = ImageBuilder(gcs_base=GCS_BASE, target_image='')
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func)
golden = '''\
def sample_component_func(a: str, b: int) -> float:
result = 3.45
if a == "succ":
result = float(b + 5)
return result
def wrapper_sample_component_func(a,b,_output_file):
output = sample_component_func(str(a),int(b))
from pathlib import Path
Path(_output_file).parent.mkdir(parents=True, exist_ok=True)
Path(_output_file).write_text(str(output))
import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("a", type=str)
parser.add_argument("b", type=int)
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())
if __name__ == "__main__":
wrapper_sample_component_func(**args)
'''
self.assertEqual(golden, generated_codes)
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func_two)
golden = '''\
def sample_component_func_two(a: str, b: int) -> float:
result = 3.45
if a == 'succ':
result = float(b + 5)
return result
def wrapper_sample_component_func_two(a,b,_output_file):
output = sample_component_func_two(str(a),int(b))
from pathlib import Path
Path(_output_file).parent.mkdir(parents=True, exist_ok=True)
Path(_output_file).write_text(str(output))
import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("a", type=str)
parser.add_argument("b", type=int)
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())
if __name__ == "__main__":
wrapper_sample_component_func_two(**args)
'''
self.assertEqual(golden, generated_codes)
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func_three)
golden = '''\
def sample_component_func_three() -> float:
return 1.0
def wrapper_sample_component_func_three(_output_file):
output = sample_component_func_three()
from pathlib import Path
Path(_output_file).parent.mkdir(parents=True, exist_ok=True)
Path(_output_file).write_text(str(output))
import argparse
parser = argparse.ArgumentParser(description="Parsing arguments")
parser.add_argument("_output_file", type=str)
args = vars(parser.parse_args())
if __name__ == "__main__":
wrapper_sample_component_func_three(**args)
'''
self.assertEqual(golden, generated_codes)
def test_generate_entrypoint(self):
""" Test entrypoint generation """
# prepare
test_data_dir = os.path.join(os.path.dirname(__file__), 'testdata')
# check
builder = ImageBuilder(gcs_base=GCS_BASE, target_image='')
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func)
golden = '''\
def sample_component_func(a: str, b: int) -> float:
result = 3.45
if a == "succ":
result = float(b + 5)
return result
def wrapper_sample_component_func(a,b):
output = sample_component_func(str(a),int(b))
with open("/output.txt", "w") as f:
f.write(str(output))
import fire
if __name__ == "__main__":
fire.Fire(wrapper_sample_component_func)
'''
self.assertEqual(golden, generated_codes)
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func_two)
golden = '''\
def sample_component_func_two(a: str, b: int) -> float:
result = 3.45
if a == 'succ':
result = float(b + 5)
return result
def wrapper_sample_component_func_two(a,b):
output = sample_component_func_two(str(a),int(b))
with open("/output.txt", "w") as f:
f.write(str(output))
import fire
if __name__ == "__main__":
fire.Fire(wrapper_sample_component_func_two)
'''
self.assertEqual(golden, generated_codes)
generated_codes = builder._generate_entrypoint(
component_func=sample_component_func_three)
golden = '''\
def sample_component_func_three() -> float:
return 1.0
def wrapper_sample_component_func_three():
output = sample_component_func_three()
with open("/output.txt", "w") as f:
f.write(str(output))
import fire
if __name__ == "__main__":
fire.Fire(wrapper_sample_component_func_three)
'''
self.assertEqual(golden, generated_codes)