@@ -486,7 +486,7 @@ def __eq__(self, other):
486
486
)
487
487
def test_text_generator_retry_success (
488
488
session ,
489
- model_name ,
489
+ model_class ,
490
490
options ,
491
491
bqml_gemini_text_generator : llm .GeminiTextGenerator ,
492
492
bqml_claude3_text_generator : llm .Claude3TextGenerator ,
@@ -567,7 +567,7 @@ def test_text_generator_retry_success(
567
567
568
568
text_generator_model = (
569
569
bqml_gemini_text_generator
570
- if (model_name == llm .GeminiTextGenerator )
570
+ if (model_class == llm .GeminiTextGenerator )
571
571
else bqml_claude3_text_generator
572
572
)
573
573
text_generator_model ._bqml_model = mock_bqml_model
@@ -602,7 +602,7 @@ def test_text_generator_retry_success(
602
602
603
603
@pytest .mark .parametrize (
604
604
(
605
- "model_name " ,
605
+ "model_class " ,
606
606
"options" ,
607
607
),
608
608
[
@@ -630,7 +630,7 @@ def test_text_generator_retry_success(
630
630
)
631
631
def test_text_generator_retry_no_progress (
632
632
session ,
633
- model_name ,
633
+ model_class ,
634
634
options ,
635
635
bqml_gemini_text_generator : llm .GeminiTextGenerator ,
636
636
bqml_claude3_text_generator : llm .Claude3TextGenerator ,
@@ -690,7 +690,7 @@ def test_text_generator_retry_no_progress(
690
690
691
691
text_generator_model = (
692
692
bqml_gemini_text_generator
693
- if (model_name == llm .GeminiTextGenerator )
693
+ if (model_class == llm .GeminiTextGenerator )
694
694
else bqml_claude3_text_generator
695
695
)
696
696
text_generator_model ._bqml_model = mock_bqml_model
0 commit comments