@@ -81,8 +81,6 @@ def test_readme_generation_str_example(
81
81
):
82
82
th = TrussHandle (custom_model_truss_dir_with_pre_and_post_str_example )
83
83
readme_contents = th .generate_readme ()
84
- with open ("tmp.md" , "w" ) as f :
85
- f .write (readme_contents )
86
84
readme_contents = readme_contents .replace ("\n " , "" )
87
85
correct_readme_contents = _read_readme ("readme_str_example.md" )
88
86
assert readme_contents == correct_readme_contents
@@ -578,8 +576,8 @@ def test_control_truss_apply_patch(custom_model_control):
578
576
running_hash = th .truss_hash_on_serving_container ()
579
577
new_model_code = """
580
578
class Model:
581
- def predict(self, request ):
582
- return [2 for i in request['inputs'] ]
579
+ def predict(self, model_input ):
580
+ return [2 for i in model_input ]
583
581
"""
584
582
patch_request = {
585
583
"hash" : "dummy" ,
@@ -620,8 +618,8 @@ def test_regular_truss_local_update_flow(custom_model_truss_dir):
620
618
model_code_file .write (
621
619
"""
622
620
class Model:
623
- def predict(self, request ):
624
- return [2 for i in request['inputs'] ]
621
+ def predict(self, model_input ):
622
+ return [2 for i in model_input ]
625
623
"""
626
624
)
627
625
result = th .docker_predict ([1 ], tag = tag )
@@ -681,8 +679,8 @@ def test_control_truss_local_update_flow(binary, python_version, custom_model_co
681
679
def predict_with_updated_model_code ():
682
680
new_model_code = """
683
681
class Model:
684
- def predict(self, request ):
685
- return [2 for i in request['inputs'] ]
682
+ def predict(self, model_input ):
683
+ return [2 for i in model_input ]
686
684
"""
687
685
model_code_file_path = custom_model_control / "model" / "model.py"
688
686
with model_code_file_path .open ("w" ) as model_code_file :
@@ -803,8 +801,8 @@ def malformed
803
801
# Should be able to fix code after
804
802
good_model_code = """
805
803
class Model:
806
- def predict(self, request ):
807
- return [2 for i in request['inputs'] ]
804
+ def predict(self, model_input ):
805
+ return [2 for i in model_input ]
808
806
"""
809
807
with model_code_file_path .open ("w" ) as model_code_file :
810
808
model_code_file .write (good_model_code )
0 commit comments