ferruzzi commented on code in PR #38693:
URL: https://github.com/apache/airflow/pull/38693#discussion_r1552580948
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -18,42 +18,153 @@
from __future__ import annotations
import json
-from typing import Generator
+from typing import TYPE_CHECKING, Generator
from unittest import mock
import pytest
from moto import mock_aws
-from airflow.providers.amazon.aws.hooks.bedrock import BedrockRuntimeHook
-from airflow.providers.amazon.aws.operators.bedrock import
BedrockInvokeModelOperator
-
-MODEL_ID = "meta.llama2-13b-chat-v1"
-PROMPT = "A very important question."
-GENERATED_RESPONSE = "An important answer."
-MOCK_RESPONSE = json.dumps(
- {
- "generation": GENERATED_RESPONSE,
- "prompt_token_count": len(PROMPT),
- "generation_token_count": len(GENERATED_RESPONSE),
- "stop_reason": "stop",
- }
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook,
BedrockRuntimeHook
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCustomizeModelOperator,
+ BedrockInvokeModelOperator,
)
-
-@pytest.fixture
-def runtime_hook() -> Generator[BedrockRuntimeHook, None, None]:
- with mock_aws():
- yield BedrockRuntimeHook(aws_conn_id="aws_default")
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
class TestBedrockInvokeModelOperator:
- @mock.patch.object(BedrockRuntimeHook, "conn")
- def test_invoke_model_prompt_good_combinations(self, mock_conn):
- mock_conn.invoke_model.return_value["body"].read.return_value =
MOCK_RESPONSE
+ def setup_method(self):
+ self.model_id = "meta.llama2-13b-chat-v1"
+ self.prompt = "A very important question."
+ self.generated_response = "An important answer."
Review Comment:
(Here and above) Yeah, it felt odd having them just hanging there outside of
any methods. I'll move them to class-level fields.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
