Skip to content

Commit 9fa3c68

Browse files
sararobcopybara-github
authored andcommitted
chore: fix cloud profiler unit test
PiperOrigin-RevId: 518014005
1 parent bb92380 commit 9fa3c68

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

tests/unit/aiplatform/test_cloud_profiler.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ def tf_profile_plugin_mock():
9292
tensorboard_plugin_profile.profile_plugin.ProfilePlugin, "capture_route"
9393
) as profile_mock:
9494
profile_mock.return_value = (
95-
wrappers.BaseResponse(
95+
wrappers.Response(
9696
json.dumps({"error": "some error"}),
9797
content_type="application/json",
9898
status=200,
@@ -341,7 +341,7 @@ def test_correct_response(self):
341341
res_dict = {"response": "OK"}
342342

343343
def my_callable(var1, var2):
344-
return wrappers.BaseResponse(
344+
return wrappers.Response(
345345
json.dumps(res_dict), content_type="application/json", status=200
346346
)
347347

0 commit comments

Comments
 (0)