Skip to content

Commit f06e804

Browse files
committed
Fixed use of validate_environment in tests
1 parent c353979 commit f06e804

File tree

1 file changed

+11
-9
lines changed

1 file changed

+11
-9
lines changed

libs/genai/tests/integration_tests/test_chat_models.py

+11-9
Original file line numberDiff line numberDiff line change
@@ -517,8 +517,8 @@ async def model_astream(context: str) -> List[BaseMessageChunk]:
517517

518518
def test_output_matches_prompt_keys() -> None:
519519
"""
520-
Validate that when response_mime_type="application/json" is specified,
521-
the output is a valid JSON format and contains the expected structure
520+
Validate that when response_mime_type="application/json" is specified,
521+
the output is a valid JSON format and contains the expected structure
522522
based on the prompt.
523523
"""
524524

@@ -554,8 +554,9 @@ def test_output_matches_prompt_keys() -> None:
554554
pytest.fail(f"Response is not valid JSON: {e}")
555555

556556
list_key = prompt_key_names["list_key"]
557-
assert list_key in response_data, \
558-
f"Expected key '{list_key}' is missing in the response."
557+
assert (
558+
list_key in response_data
559+
), f"Expected key '{list_key}' is missing in the response."
559560
grocery_items = response_data[list_key]
560561
assert isinstance(grocery_items, list), f"'{list_key}' should be a list."
561562

@@ -569,21 +570,22 @@ def test_output_matches_prompt_keys() -> None:
569570

570571
print("Response matches the key names specified in the prompt.")
571572

573+
572574
def test_validate_response_mime_type_and_schema() -> None:
573575
"""
574576
Test that `response_mime_type` and `response_schema` are validated correctly.
575577
Ensure valid combinations of `response_mime_type` and `response_schema` pass,
576578
and invalid ones raise appropriate errors.
577579
"""
578580

579-
valid_model = ChatGoogleGenerativeAI(
581+
llm = ChatGoogleGenerativeAI(
580582
model="gemini-1.5-pro",
581583
response_mime_type="application/json",
582584
response_schema={"type": "list", "items": {"type": "object"}}, # Example schema
583585
)
584586

585587
try:
586-
valid_model.validate_environment()
588+
llm.invoke("Hello")
587589
except ValueError as e:
588590
pytest.fail(f"Validation failed unexpectedly with valid parameters: {e}")
589591

@@ -592,14 +594,14 @@ def test_validate_response_mime_type_and_schema() -> None:
592594
model="gemini-1.5-pro",
593595
response_mime_type="invalid/type",
594596
response_schema={"type": "list", "items": {"type": "object"}},
595-
).validate_environment()
597+
).invoke("Hello")
596598

597599
try:
598600
ChatGoogleGenerativeAI(
599601
model="gemini-1.5-pro",
600602
response_mime_type="application/json",
601-
).validate_environment()
603+
).invoke("Hello")
602604
except ValueError as e:
603605
pytest.fail(
604606
f"Validation failed unexpectedly with a valid MIME type and no schema: {e}"
605-
)
607+
)

0 commit comments

Comments
 (0)