Skip to content

Commit 7dc0a4a

Browse files
Updated test case assertions for consistency.
1 parent cdd8648 commit 7dc0a4a

File tree

1 file changed

+21
-41
lines changed

1 file changed

+21
-41
lines changed

model_armor/snippets/snippets_test.py

Lines changed: 21 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,7 @@ def test_create_template(
447447
project_id: str, location_id: str, template_id: str
448448
) -> None:
449449
template = create_model_armor_template(project_id, location_id, template_id)
450-
assert template is not None
450+
assert template
451451

452452

453453
def test_get_template(
@@ -509,17 +509,16 @@ def test_create_model_armor_template_with_basic_sdp(
509509
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
510510
)
511511

512-
assert (
513-
created_template.name == expected_name_format
514-
), "Template name does not match the expected format."
512+
assert created_template.name == expected_name_format
515513

516514
filter_enforcement = (
517515
created_template.filter_config.sdp_settings.basic_config.filter_enforcement
518516
)
519517

520518
assert (
521-
filter_enforcement.name == "ENABLED"
522-
), f"Expected filter_enforcement to be ENABLED, but got {filter_enforcement}"
519+
filter_enforcement.name
520+
== modelarmor_v1.SdpBasicConfig.filter_enforcement.ENABLED.name
521+
)
523522

524523

525524
def test_create_model_armor_template_with_advanced_sdp(
@@ -533,33 +532,27 @@ def test_create_model_armor_template_with_advanced_sdp(
533532
that matches the expected format.
534533
"""
535534

536-
sdp_inspect_template_id, sdr_deidentify_template_id = sdp_templates
535+
sdp_inspect_template_id, sdp_deidentify_template_id = sdp_templates
537536
created_template = create_model_armor_template_with_advanced_sdp(
538537
project_id,
539538
location_id,
540539
template_id,
541540
sdp_inspect_template_id,
542-
sdr_deidentify_template_id,
541+
sdp_deidentify_template_id,
543542
)
544543

545544
expected_name_format = (
546545
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
547546
)
548547

549-
assert (
550-
created_template.name == expected_name_format
551-
), "Template name does not match the expected format."
548+
assert created_template.name == expected_name_format
552549

553550
advanced_config = (
554551
created_template.filter_config.sdp_settings.advanced_config
555552
)
556-
assert (
557-
advanced_config.inspect_template == sdp_inspect_template_id
558-
), f"Expected inspect_template to be {sdp_inspect_template_id}, but got {advanced_config.inspect_template}"
553+
assert advanced_config.inspect_template == sdp_inspect_template_id
559554

560-
assert (
561-
advanced_config.deidentify_template == sdr_deidentify_template_id
562-
), f"Expected deidentify_template to be {sdr_deidentify_template_id}, but got {advanced_config.deidentify_template}"
555+
assert advanced_config.deidentify_template == sdp_deidentify_template_id
563556

564557

565558
def test_create_model_armor_template_with_metadata(
@@ -578,9 +571,7 @@ def test_create_model_armor_template_with_metadata(
578571
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
579572
)
580573

581-
assert (
582-
created_template.name == expected_name_format
583-
), "Template name does not match the expected format."
574+
assert created_template.name == expected_name_format
584575
assert created_template.template_metadata.ignore_partial_invocation_failures
585576
assert created_template.template_metadata.log_sanitize_operations
586577

@@ -601,18 +592,14 @@ def test_create_model_armor_template_with_labels(
601592
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
602593
)
603594

604-
assert (
605-
created_template.name == expected_name_format
606-
), "Template name does not match the expected format."
595+
assert created_template.name == expected_name_format
607596

608597
template_with_labels = get_model_armor_template(
609598
project_id, location_id, template_id
610599
)
611600

612601
for key, value in expected_labels.items():
613-
assert (
614-
template_with_labels.labels.get(key) == value
615-
), f"Label {key} does not match. Expected: {value}, Got: {template_with_labels.labels.get(key)}"
602+
assert template_with_labels.labels.get(key) == value
616603

617604

618605
def test_list_model_armor_templates_with_filter(
@@ -636,7 +623,7 @@ def test_list_model_armor_templates_with_filter(
636623

637624
assert any(
638625
template.name == expected_template_name for template in templates
639-
), "Template does not exist in the list"
626+
)
640627

641628

642629
def test_update_model_armor_template_metadata(
@@ -658,9 +645,7 @@ def test_update_model_armor_template_metadata(
658645
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
659646
)
660647

661-
assert (
662-
updated_template.name == expected_name_format
663-
), "Template name does not match the expected format."
648+
assert updated_template.name == expected_name_format
664649
assert updated_template.template_metadata.ignore_partial_invocation_failures
665650
assert updated_template.template_metadata.log_sanitize_operations
666651

@@ -685,18 +670,14 @@ def test_update_model_armor_template_labels(
685670
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
686671
)
687672

688-
assert (
689-
updated_template.name == expected_name_format
690-
), "Template name does not match the expected format."
673+
assert updated_template.name == expected_name_format
691674

692675
template_with_lables = get_model_armor_template(
693676
project_id, location_id, template_id
694677
)
695678

696679
for key, value in expected_labels.items():
697-
assert (
698-
template_with_lables.labels.get(key) == value
699-
), f"Label {key} does not match. Expected: {value}, Got: {template_with_lables.labels.get(key)}"
680+
assert template_with_lables.labels.get(key) == value
700681

701682

702683
def test_update_model_armor_template_with_mask_configuration(
@@ -718,16 +699,15 @@ def test_update_model_armor_template_with_mask_configuration(
718699
f"projects/{project_id}/locations/{location_id}/templates/{template_id}"
719700
)
720701

721-
assert (
722-
updated_template.name == expected_name_format
723-
), "Template name does not match the expected format."
702+
assert updated_template.name == expected_name_format
724703

725704
filter_enforcement = (
726705
updated_template.filter_config.sdp_settings.basic_config.filter_enforcement
727706
)
728707
assert (
729-
filter_enforcement.name != "ENABLED"
730-
), f"Expected filter_enforcement not to be ENABLED, but got {filter_enforcement}"
708+
filter_enforcement.name
709+
!= modelarmor_v1.SdpBasicConfig.filter_enforcement.ENABLED.name
710+
)
731711

732712

733713
def test_sanitize_user_prompt_with_all_rai_filter_template(

0 commit comments

Comments
 (0)