From e8fbe1cb016e9524896e78fca25ce8aed6771676 Mon Sep 17 00:00:00 2001 From: Thiago Castro Ferreira Date: Fri, 22 Nov 2024 16:03:14 -0300 Subject: [PATCH] Default parameter on ModelResponse.get --- aixplain/modules/model/response.py | 7 +++++-- tests/unit/model_test.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/aixplain/modules/model/response.py b/aixplain/modules/model/response.py index 94ddcb9d..99e32074 100644 --- a/aixplain/modules/model/response.py +++ b/aixplain/modules/model/response.py @@ -40,8 +40,11 @@ def __getitem__(self, key: Text) -> Any: return self.run_time raise KeyError(f"Key '{key}' not found in ModelResponse.") - def get(self, key: Text) -> Any: - return self[key] + def get(self, key: Text, default: Optional[Any] = None) -> Any: + try: + return self[key] + except KeyError: + return default def __repr__(self) -> str: fields = [] diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py index b45b6ae0..a2463a8d 100644 --- a/tests/unit/model_test.py +++ b/tests/unit/model_test.py @@ -87,7 +87,7 @@ def test_success_poll(): hyp_response = test_model.poll(poll_url=poll_url) assert isinstance(hyp_response, ModelResponse) assert hyp_response["completed"] == ref_response["completed"] - assert hyp_response["status"] == ResponseStatus.SUCCESS + assert hyp_response.get("status") == ResponseStatus.SUCCESS def test_failed_poll():