ferruzzi commented on code in PR #39500:
URL: https://github.com/apache/airflow/pull/39500#discussion_r1596065112
##########
tests/providers/amazon/aws/operators/test_bedrock.py:
##########
@@ -346,3 +347,168 @@ def test_id_returned(self, mock_conn):
result = self.operator.execute({})
assert result == self.INGESTION_JOB_ID
+
+
+class TestBedrockRaGOperator:
+ VECTOR_SEARCH_CONFIG = {"filter": {"equals": {"key": "some key", "value":
"some value"}}}
+ KNOWLEDGE_BASE_ID = "knowledge_base_id"
+ SOURCES = [{"sourceType": "S3", "s3Location": "bucket"}]
+ MODEL_ARN = "model arn"
+
+ @pytest.mark.parametrize(
+ "source_type, vector_search_config, knowledge_base_id, sources,
expect_success",
+ [
+ pytest.param(
+ "invalid_source_type",
+ None,
+ None,
+ None,
+ False,
+ id="invalid_source_type",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ None,
+ None,
+ False,
+ id="KNOWLEDGE_BASE_without_knowledge_base_id_fails",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ KNOWLEDGE_BASE_ID,
+ SOURCES,
+ False,
+ id="KNOWLEDGE_BASE_with_sources_fails",
+ ),
+ pytest.param(
+ "KNOWLEDGE_BASE",
+ VECTOR_SEARCH_CONFIG,
+ KNOWLEDGE_BASE_ID,
+ None,
+ True,
+ id="KNOWLEDGE_BASE_passes",
+ ),
Review Comment:
Sure. vector config is optional with a knowledge base, but forbidden when
using external. I can add that check.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
