fixing tests

mckinsey · Aug 8, 2024 · 768bead · 768bead
1 parent ea4ffc7
commit 768bead
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/vizro-ai/tests/unit/vizro-ai/dashboard/test_pydantic_output.py b/vizro-ai/tests/unit/vizro-ai/dashboard/test_pydantic_output.py
@@ -1,15 +1,14 @@
 import pytest
 import vizro.models as vm
-from tests_utils.asserts import assert_component_equal
 from vizro_ai.dashboard._pydantic_output import _create_message_content, _create_prompt_template, _get_pydantic_model
 
 
 def test_get_pydantic_model_valid(component_description, fake_llm):
     pydantic_output = _get_pydantic_model(
         query=component_description, llm_model=fake_llm, response_model=vm.Card, df_info=None
     )
-    assert_component_equal(
-        pydantic_output.build(), vm.Card(text="this is a card", href="").build(), keys_to_strip={"id"}
+    assert pydantic_output.dict(exclude={"id": True}) == vm.Card(text="this is a card", href="").dict(
+        exclude={"id": True}
     )