chore: adding google3 testing support for model-builder SDK sample tests

PiperOrigin-RevId: 507602651

samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample.py

@@ -20,5 +20,8 @@ def assign_artifact_as_execution_input_sample(
     execution: aiplatform.Execution,
     artifact: aiplatform.Artifact,
 ):
+
     execution.assign_input_artifacts([artifact])
+
+
 #  [END aiplatform_sdk_assign_artifact_as_execution_input_sample]

samples/model-builder/experiment_tracking/assign_artifact_as_execution_input_sample_test.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import assign_artifact_as_execution_input_sample
 from google.cloud import aiplatform
 
+from experiment_tracking import assign_artifact_as_execution_input_sample
+
 
 def test_assign_artifact_as_execution_input_sample(
     mock_get_execution,

samples/model-builder/experiment_tracking/assign_artifact_as_execution_output_sample_test.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import assign_artifact_as_execution_output_sample
 from google.cloud import aiplatform
 
+from experiment_tracking import assign_artifact_as_execution_output_sample
+
 
 def test_assign_artifact_as_execution_output_sample(
     mock_get_execution,

samples/model-builder/experiment_tracking/create_artifact_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import create_artifact_sample
-
+from experiment_tracking import create_artifact_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import create_artifact_with_sdk_sample
-
+from experiment_tracking import create_artifact_with_sdk_sample
 import test_constants as constants
 
 
-def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_schema_base_artifact):
+def test_create_artifact_with_sdk_sample(
+    mock_artifact, mock_create_schema_base_artifact
+):
     artifact = create_artifact_with_sdk_sample.create_artifact_sample(
         project=constants.PROJECT,
         location=constants.LOCATION,
@@ -30,6 +31,7 @@ def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_schema_base_
     )
 
     mock_create_schema_base_artifact.assert_called_with(
-        project='abc', location='us-central1')
+        project="abc", location="us-central1"
+    )
 
     assert artifact is mock_artifact

samples/model-builder/experiment_tracking/create_context_with_sdk_sample_test.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import create_context_with_sdk_sample
-
+from experiment_tracking import create_context_with_sdk_sample
 import test_constants as constants
 
 
 def test_create_context_sample(
-    mock_sdk_init, mock_create_schema_base_context, mock_context,
+    mock_sdk_init,
+    mock_create_schema_base_context,
+    mock_context,
 ):
     exc = create_context_with_sdk_sample.create_context_sample(
         display_name=constants.DISPLAY_NAME,
@@ -31,7 +32,8 @@ def test_create_context_sample(
     )
 
     mock_sdk_init.assert_called_with(
-        project=constants.PROJECT, location=constants.LOCATION,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
     )
 
     mock_create_schema_base_context.assert_called_with()

samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import create_execution_with_sdk_sample
-
+from experiment_tracking import create_execution_with_sdk_sample
 import test_constants as constants
 
 
 def test_create_execution_sample(
-    mock_sdk_init, mock_create_artifact, mock_create_schema_base_execution, mock_execution,
+    mock_sdk_init,
+    mock_create_artifact,
+    mock_create_schema_base_execution,
+    mock_execution,
 ):
 
     input_art = mock_create_artifact()
@@ -37,7 +39,8 @@ def test_create_execution_sample(
     )
 
     mock_sdk_init.assert_called_with(
-        project=constants.PROJECT, location=constants.LOCATION,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
     )
 
     mock_create_schema_base_execution.assert_called_with()

samples/model-builder/experiment_tracking/create_experiment_run_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import create_experiment_run_sample
-
+from experiment_tracking import create_experiment_run_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/create_experiment_sample_test.py

@@ -12,9 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-import create_experiment_sample
-
+from experiment_tracking import create_experiment_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/delete_artifact_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import delete_artifact_sample
-
+from experiment_tracking import delete_artifact_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/delete_context_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import delete_context_sample
-
+from experiment_tracking import delete_context_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/delete_execution_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import delete_execution_sample
-
+from experiment_tracking import delete_execution_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/delete_experiment_run_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import delete_experiment_run_sample
-
+from experiment_tracking import delete_experiment_run_sample
 import test_constants
 
 

samples/model-builder/experiment_tracking/delete_experiment_sample_test.py

@@ -12,9 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-import delete_experiment_sample
-
+from experiment_tracking import delete_experiment_sample
 import test_constants
 
 

samples/model-builder/experiment_tracking/end_experiment_run_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import end_experiment_run_sample
-
+from experiment_tracking import end_experiment_run_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_artifact_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_artifact_sample
-
+from experiment_tracking import get_artifact_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_artifact_with_uri_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_artifact_with_uri_sample
-
+from experiment_tracking import get_artifact_with_uri_sample
 import test_constants
 
 

samples/model-builder/experiment_tracking/get_context_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_context_sample
-
+from experiment_tracking import get_context_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_execution_input_artifacts_sample_test.py

@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_execution_input_artifacts_sample
 from google.cloud import aiplatform
 
+from experiment_tracking import get_execution_input_artifacts_sample
 
-def test_aiplatform_sdk_get_execution_input_artifacts_sample(
-    mock_get_execution
-):
+
+def test_aiplatform_sdk_get_execution_input_artifacts_sample(mock_get_execution):
     exc = aiplatform.Execution()
     get_execution_input_artifacts_sample.get_execution_input_artifacts_sample(
         execution=exc

samples/model-builder/experiment_tracking/get_execution_output_artifacts_sample_test.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_execution_output_artifacts_sample
+
 from google.cloud import aiplatform
 
+from experiment_tracking import get_execution_output_artifacts_sample
+
 
-def test_aiplatform_sdk_get_execution_output_artifacts_sample(
-    mock_get_execution
-):
+def test_aiplatform_sdk_get_execution_output_artifacts_sample(mock_get_execution):
     exc = aiplatform.Execution()
     get_execution_output_artifacts_sample.get_execution_output_artifacts_sample(
         execution=exc

samples/model-builder/experiment_tracking/get_execution_sample_test.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_execution_sample
-
+from experiment_tracking import get_execution_sample
 import test_constants
 
 

samples/model-builder/experiment_tracking/get_experiment_data_frame_sample_test.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_data_frame_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_data_frame_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_experiment_model_sample_test.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_model_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_model_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_experiment_run_artifacts_sample_test.py

@@ -12,23 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_run_artifacts_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_run_artifacts_sample
 import test_constants as constants
 
 
 @pytest.mark.usefixtures("mock_get_run")
 def test_get_experiment_run_artifact_sample(mock_get_artifacts, mock_artifacts):
 
-    artifacts = (
-        get_experiment_run_artifacts_sample.get_experiment_run_artifacts_sample(
-            run_name=constants.EXPERIMENT_RUN_NAME,
-            experiment=constants.EXPERIMENT_NAME,
-            project=constants.PROJECT,
-            location=constants.LOCATION,
-        )
+    artifacts = get_experiment_run_artifacts_sample.get_experiment_run_artifacts_sample(
+        run_name=constants.EXPERIMENT_RUN_NAME,
+        experiment=constants.EXPERIMENT_NAME,
+        project=constants.PROJECT,
+        location=constants.LOCATION,
     )
 
     mock_get_artifacts.assert_called_with()

samples/model-builder/experiment_tracking/get_experiment_run_classification_metrics_sample_test.py

@@ -12,15 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_run_classification_metrics_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_run_classification_metrics_sample
 import test_constants as constants
 
 
 @pytest.mark.usefixtures("mock_get_run")
-def test_get_experiment_run_classification_metrics_sample(mock_get_classification_metrics, mock_classification_metrics):
+def test_get_experiment_run_classification_metrics_sample(
+    mock_get_classification_metrics, mock_classification_metrics
+):
 
     classification_metrics = get_experiment_run_classification_metrics_sample.get_experiment_run_classification_metrics_sample(
         run_name=constants.EXPERIMENT_RUN_NAME,

samples/model-builder/experiment_tracking/get_experiment_run_metrics_sample_test.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_run_metrics_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_run_metrics_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_experiment_run_models_test.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_run_models_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_run_models_sample
 import test_constants as constants
 
 

samples/model-builder/experiment_tracking/get_experiment_run_params_sample_test.py

@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import get_experiment_run_params_sample
-
 import pytest
 
+from experiment_tracking import get_experiment_run_params_sample
 import test_constants as constants