diff --git a/src/neo4j_graphrag/experimental/pipeline/pipeline.py b/src/neo4j_graphrag/experimental/pipeline/pipeline.py index 3d004eb8..5edc2783 100644 --- a/src/neo4j_graphrag/experimental/pipeline/pipeline.py +++ b/src/neo4j_graphrag/experimental/pipeline/pipeline.py @@ -15,12 +15,12 @@ from __future__ import annotations import asyncio +import datetime import enum import logging import uuid import warnings from collections import defaultdict -from datetime import datetime from timeit import default_timer from typing import Any, AsyncGenerator, Optional @@ -61,7 +61,9 @@ class RunStatus(enum.Enum): class RunResult(BaseModel): status: RunStatus = RunStatus.DONE result: Optional[DataModel] = None - timestamp: datetime = Field(default_factory=datetime.utcnow) + timestamp: datetime.datetime = Field( + default_factory=lambda: datetime.datetime.now(datetime.timezone.utc) + ) class TaskPipelineNode(PipelineNode): diff --git a/tests/unit/experimental/pipeline/test_pipeline.py b/tests/unit/experimental/pipeline/test_pipeline.py index 8191da61..7eeae9cc 100644 --- a/tests/unit/experimental/pipeline/test_pipeline.py +++ b/tests/unit/experimental/pipeline/test_pipeline.py @@ -16,12 +16,14 @@ import asyncio import tempfile +from typing import Sized from unittest import mock from unittest.mock import AsyncMock, call, patch import pytest from neo4j_graphrag.experimental.pipeline import Component, Pipeline from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError +from neo4j_graphrag.experimental.pipeline.pipeline import RunResult from .components import ( ComponentAdd, @@ -404,3 +406,8 @@ def test_pipeline_draw_missing_pygraphviz_dep() -> None: t = tempfile.NamedTemporaryFile() with pytest.raises(ImportError): pipe.draw(t.name) + + +def test_run_result_no_warning(recwarn: Sized) -> None: + RunResult() + assert len(recwarn) == 0