@@ -113,85 +113,69 @@ def test_create_lineage_when_no_lineage_exists_with_fg_only():
113113 transformation_code = TRANSFORMATION_CODE_INPUT_1 ,
114114 sagemaker_session = SAGEMAKER_SESSION_MOCK ,
115115 )
116- with (
117- patch .object (
118- FeatureGroupLineageEntityHandler ,
119- "retrieve_feature_group_context_arns" ,
120- side_effect = [
121- FEATURE_GROUP_INPUT [0 ],
122- FEATURE_GROUP_INPUT [1 ],
123- FEATURE_GROUP_INPUT [0 ],
124- ],
125- ) as retrieve_feature_group_context_arns_method
126- patch .object (
127- S3LineageEntityHandler ,
128- "retrieve_raw_data_artifact" ,
129- side_effect = [
130- RAW_DATA_INPUT_ARTIFACTS [0 ],
131- RAW_DATA_INPUT_ARTIFACTS [1 ],
132- RAW_DATA_INPUT_ARTIFACTS [2 ],
133- RAW_DATA_INPUT_ARTIFACTS [3 ],
134- ],
135- ) as retrieve_raw_data_artifact_method ,
136- patch .object (
137- S3LineageEntityHandler ,
138- "create_transformation_code_artifact" ,
139- return_value = TRANSFORMATION_CODE_ARTIFACT_1 ,
140- ) as create_transformation_code_artifact_method ,
141- patch .object (
142- PipelineLineageEntityHandler ,
143- "load_pipeline_context" ,
144- side_effect = RESOURCE_NOT_FOUND_EXCEPTION ,
145- ) as load_pipeline_context_method ,
146- patch .object (
147- PipelineLineageEntityHandler ,
148- "create_pipeline_context" ,
149- return_value = PIPELINE_CONTEXT ,
150- ),
151- patch .object (
152- PipelineVersionLineageEntityHandler ,
153- "create_pipeline_version_context" ,
154- return_value = PIPELINE_VERSION_CONTEXT ,
155- ),
156- patch .object (
157- PipelineVersionLineageEntityHandler ,
158- "load_pipeline_version_context" ,
159- return_value = PIPELINE_VERSION_CONTEXT ,
160- ) as load_pipeline_version_context_method ,
161- patch .object (
162- LineageAssociationHandler ,
163- "list_upstream_associations" ,
164- side_effect = [
165- generate_pipeline_version_upstream_feature_group_list (),
166- [],
167- generate_pipeline_version_upstream_transformation_code (),
168- ],
169- ) as list_upstream_associations_method ,
170- patch .object (
171- LineageAssociationHandler ,
172- "list_downstream_associations" ,
173- return_value = generate_pipeline_version_downstream_feature_group (),
174- ) as list_downstream_associations_method ,
175- patch .object (
176- PipelineLineageEntityHandler ,
177- "update_pipeline_context" ,
178- ) as update_pipeline_context_method ,
179- patch .object (
180- LineageAssociationHandler , "add_upstream_feature_group_data_associations"
181- ) as add_upstream_feature_group_data_associations_method ,
182- patch .object (
183- LineageAssociationHandler , "add_downstream_feature_group_data_associations"
184- ) as add_downstream_feature_group_data_associations_method ,
185- patch .object (
186- LineageAssociationHandler , "add_upstream_raw_data_associations"
187- ) as add_upstream_raw_data_associations_method ,
188- patch .object (
189- LineageAssociationHandler , "add_upstream_transformation_code_associations"
190- ) as add_upstream_transformation_code_associations_method ,
191- patch .object (
192- LineageAssociationHandler , "add_pipeline_and_pipeline_version_association"
193- ) as add_pipeline_and_pipeline_version_association_method ,
194- ):
116+ with patch .object (
117+ FeatureGroupLineageEntityHandler ,
118+ "retrieve_feature_group_context_arns" ,
119+ side_effect = [
120+ FEATURE_GROUP_INPUT [0 ],
121+ FEATURE_GROUP_INPUT [1 ],
122+ FEATURE_GROUP_INPUT [0 ],
123+ ],
124+ ) as retrieve_feature_group_context_arns_method , patch .object (
125+ S3LineageEntityHandler ,
126+ "retrieve_raw_data_artifact" ,
127+ side_effect = [
128+ RAW_DATA_INPUT_ARTIFACTS [0 ],
129+ RAW_DATA_INPUT_ARTIFACTS [1 ],
130+ RAW_DATA_INPUT_ARTIFACTS [2 ],
131+ RAW_DATA_INPUT_ARTIFACTS [3 ],
132+ ],
133+ ) as retrieve_raw_data_artifact_method , patch .object (
134+ S3LineageEntityHandler ,
135+ "create_transformation_code_artifact" ,
136+ return_value = TRANSFORMATION_CODE_ARTIFACT_1 ,
137+ ) as create_transformation_code_artifact_method , patch .object (
138+ PipelineLineageEntityHandler ,
139+ "load_pipeline_context" ,
140+ side_effect = RESOURCE_NOT_FOUND_EXCEPTION ,
141+ ) as load_pipeline_context_method , patch .object (
142+ PipelineLineageEntityHandler ,
143+ "create_pipeline_context" ,
144+ return_value = PIPELINE_CONTEXT ,
145+ ), patch .object (
146+ PipelineVersionLineageEntityHandler ,
147+ "create_pipeline_version_context" ,
148+ return_value = PIPELINE_VERSION_CONTEXT ,
149+ ), patch .object (
150+ PipelineVersionLineageEntityHandler ,
151+ "load_pipeline_version_context" ,
152+ return_value = PIPELINE_VERSION_CONTEXT ,
153+ ) as load_pipeline_version_context_method , patch .object (
154+ LineageAssociationHandler ,
155+ "list_upstream_associations" ,
156+ side_effect = [
157+ generate_pipeline_version_upstream_feature_group_list (),
158+ [],
159+ generate_pipeline_version_upstream_transformation_code (),
160+ ],
161+ ) as list_upstream_associations_method , patch .object (
162+ LineageAssociationHandler ,
163+ "list_downstream_associations" ,
164+ return_value = generate_pipeline_version_downstream_feature_group (),
165+ ) as list_downstream_associations_method , patch .object (
166+ PipelineLineageEntityHandler ,
167+ "update_pipeline_context" ,
168+ ) as update_pipeline_context_method , patch .object (
169+ LineageAssociationHandler , "add_upstream_feature_group_data_associations"
170+ ) as add_upstream_feature_group_data_associations_method , patch .object (
171+ LineageAssociationHandler , "add_downstream_feature_group_data_associations"
172+ ) as add_downstream_feature_group_data_associations_method , patch .object (
173+ LineageAssociationHandler , "add_upstream_raw_data_associations"
174+ ) as add_upstream_raw_data_associations_method , patch .object (
175+ LineageAssociationHandler , "add_upstream_transformation_code_associations"
176+ ) as add_upstream_transformation_code_associations_method , patch .object (
177+ LineageAssociationHandler , "add_pipeline_and_pipeline_version_association"
178+ ) as add_pipeline_and_pipeline_version_association_method :
195179 lineage_handler .create_lineage ()
196180
197181 retrieve_feature_group_context_arns_method .assert_has_calls (
0 commit comments