Skip to content

Commit

Permalink
Fix/datetime warning (#202)
Browse files Browse the repository at this point in the history
* Fix deprecation warning (utcnow)

* Mypy
  • Loading branch information
stellasia authored Oct 25, 2024
1 parent cb96815 commit 99bf50e
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 2 deletions.
6 changes: 4 additions & 2 deletions src/neo4j_graphrag/experimental/pipeline/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down
7 changes: 7 additions & 0 deletions tests/unit/experimental/pipeline/test_pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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

0 comments on commit 99bf50e

Please sign in to comment.