@@ -527,6 +527,12 @@ def mock_artifact():
527
527
yield mock
528
528
529
529
530
+ @pytest .fixture
531
+ def mock_context ():
532
+ mock = MagicMock (aiplatform .Context )
533
+ yield mock
534
+
535
+
530
536
@pytest .fixture
531
537
def mock_experiment ():
532
538
mock = MagicMock (aiplatform .Experiment )
@@ -613,13 +619,41 @@ def mock_artifact_get(mock_artifact):
613
619
yield mock_artifact_get
614
620
615
621
622
+ @pytest .fixture
623
+ def mock_context_get (mock_context ):
624
+ with patch .object (aiplatform .Context , "get" ) as mock_context_get :
625
+ mock_context_get .return_value = mock_context
626
+ yield mock_context_get
627
+
628
+
616
629
@pytest .fixture
617
630
def mock_pipeline_job_create (mock_pipeline_job ):
618
631
with patch .object (aiplatform , "PipelineJob" ) as mock_pipeline_job_create :
619
632
mock_pipeline_job_create .return_value = mock_pipeline_job
620
633
yield mock_pipeline_job_create
621
634
622
635
636
+ @pytest .fixture
637
+ def mock_artifact_delete ():
638
+ with patch .object (aiplatform .Artifact , "delete" ) as mock_artifact_delete :
639
+ mock_artifact_delete .return_value = None
640
+ yield mock_artifact_delete
641
+
642
+
643
+ @pytest .fixture
644
+ def mock_execution_delete ():
645
+ with patch .object (aiplatform .Execution , "delete" ) as mock_execution_delete :
646
+ mock_execution_delete .return_value = None
647
+ yield mock_execution_delete
648
+
649
+
650
+ @pytest .fixture
651
+ def mock_context_delete ():
652
+ with patch .object (aiplatform .Context , "delete" ) as mock_context_delete :
653
+ mock_context_delete .return_value = None
654
+ yield mock_context_delete
655
+
656
+
623
657
@pytest .fixture
624
658
def mock_pipeline_job_submit (mock_pipeline_job ):
625
659
with patch .object (mock_pipeline_job , "submit" ) as mock_pipeline_job_submit :
0 commit comments