diff --git a/vizro-core/tests/unit/vizro/models/test_types.py b/vizro-core/tests/unit/vizro/models/test_types.py
index e093e067c..db00a89cf 100644
--- a/vizro-core/tests/unit/vizro/models/test_types.py
+++ b/vizro-core/tests/unit/vizro/models/test_types.py
@@ -4,14 +4,15 @@
 import plotly.graph_objects as go
 import plotly.io as pio
 import pytest
-from pydantic import Field, ValidationError
+from pydantic import Field, ValidationError, field_validator
+from pydantic.json_schema import SkipJsonSchema
 
 # try:
 #     from pydantic.v1 import Field, ValidationError
 # except ImportError:  # pragma: no cov
 #     from pydantic import Field, ValidationError
 from vizro.models import VizroBaseModel
-from vizro.models.types import CapturedCallable, capture
+from vizro.models.types import CapturedCallable, capture, validate_captured_callable
 
 
 def positional_only_function(a, /):
@@ -163,12 +164,18 @@ def invalid_decorated_graph_function():
 
 class ModelWithAction(VizroBaseModel):
     # The import_path here makes it possible to import the above function using getattr(import_path, _target_).
-    function: CapturedCallable = Field(..., import_path=__name__, mode="action")
+    function: SkipJsonSchema[CapturedCallable] = Field(
+        ..., json_schema_extra={"mode": "action", "import_path": __name__}
+    )
+    _validate_figure = field_validator("function", mode="before")(validate_captured_callable)
 
 
 class ModelWithGraph(VizroBaseModel):
     # The import_path here makes it possible to import the above function using getattr(import_path, _target_).
-    function: CapturedCallable = Field(..., import_path=__name__, mode="graph")
+    function: SkipJsonSchema[CapturedCallable] = Field(
+        ..., json_schema_extra={"mode": "graph", "import_path": __name__}
+    )
+    _validate_figure = field_validator("function", mode="before")(validate_captured_callable)
 
 
 class TestModelFieldPython: