This is an automated email from the ASF dual-hosted git repository.
shahar pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fabc5c35da0 Add exception test for GenAICountTokensOperator (#61391)
fabc5c35da0 is described below
commit fabc5c35da0219f904817e54ee4393f1ae07b101
Author: Darshan Gorasiya <[email protected]>
AuthorDate: Sun Mar 8 02:34:22 2026 +0000
Add exception test for GenAICountTokensOperator (#61391)
---
.../unit/google/cloud/operators/test_gen_ai.py | 25 ++++++++++++++++++++++
1 file changed, 25 insertions(+)
diff --git a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
index 761a835678c..9446f5b858a 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
@@ -238,6 +238,31 @@ class TestGenAICountTokensOperator:
config=None,
)
+ @mock.patch(GEN_AI_PATH.format("GenAIGenerativeModelHook"))
+ def test_execute_propagates_client_error(self, mock_hook):
+ """Test that GenAICountTokensOperator propagates ClientError from the
hook."""
+ # Bypass __init__ to avoid version-dependent constructor signature
differences
+ # between google-genai versions (e.g., 1.2.0 vs 1.63.0).
+ mock_hook.return_value.count_tokens.side_effect =
ClientError.__new__(ClientError)
+
+ op = GenAICountTokensOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ contents=CONTENTS,
+ model=GEMINI_MODEL,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
+ with pytest.raises(ClientError):
+ op.execute(context={"ti": mock.MagicMock()})
+
+ mock_hook.assert_called_once_with(
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ )
+
class TestGenAICreateCachedContentOperator:
@mock.patch(GEN_AI_PATH.format("GenAIGenerativeModelHook"))
