Skip to content

Commit 60061a5

Browse files
committed
fix a variable name typo
1 parent cc838e9 commit 60061a5

File tree

2 files changed

+6
-6
lines changed

2 files changed

+6
-6
lines changed

tests/system/small/ml/conftest.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -393,7 +393,7 @@ def imported_xgboost_model(
393393
@pytest.fixture(scope="session")
394394
def bqml_gemini_text_generator(bq_connection, session) -> llm.GeminiTextGenerator:
395395
return llm.GeminiTextGenerator(
396-
model_name="gemini-1.5-flash-002", # Or your desired Gemini model
396+
model_name="gemini-1.5-flash-002",
397397
connection_name=bq_connection,
398398
session=session,
399399
)

tests/system/small/ml/test_llm.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -486,7 +486,7 @@ def __eq__(self, other):
486486
)
487487
def test_text_generator_retry_success(
488488
session,
489-
model_name,
489+
model_class,
490490
options,
491491
bqml_gemini_text_generator: llm.GeminiTextGenerator,
492492
bqml_claude3_text_generator: llm.Claude3TextGenerator,
@@ -567,7 +567,7 @@ def test_text_generator_retry_success(
567567

568568
text_generator_model = (
569569
bqml_gemini_text_generator
570-
if (model_name == llm.GeminiTextGenerator)
570+
if (model_class == llm.GeminiTextGenerator)
571571
else bqml_claude3_text_generator
572572
)
573573
text_generator_model._bqml_model = mock_bqml_model
@@ -602,7 +602,7 @@ def test_text_generator_retry_success(
602602

603603
@pytest.mark.parametrize(
604604
(
605-
"model_name",
605+
"model_class",
606606
"options",
607607
),
608608
[
@@ -630,7 +630,7 @@ def test_text_generator_retry_success(
630630
)
631631
def test_text_generator_retry_no_progress(
632632
session,
633-
model_name,
633+
model_class,
634634
options,
635635
bqml_gemini_text_generator: llm.GeminiTextGenerator,
636636
bqml_claude3_text_generator: llm.Claude3TextGenerator,
@@ -690,7 +690,7 @@ def test_text_generator_retry_no_progress(
690690

691691
text_generator_model = (
692692
bqml_gemini_text_generator
693-
if (model_name == llm.GeminiTextGenerator)
693+
if (model_class == llm.GeminiTextGenerator)
694694
else bqml_claude3_text_generator
695695
)
696696
text_generator_model._bqml_model = mock_bqml_model

0 commit comments

Comments
 (0)