Incorporate nodestream-proj/nodestream#257

Author: zprobst

nodestream_plugin_neo4j/ingest_query_builder.py

@@ -1,5 +1,4 @@
 import re
-from datetime import datetime, timedelta
 from functools import cache, wraps
 from typing import Iterable
 
@@ -18,6 +17,8 @@
     TimeToLiveConfiguration,
 )
 from nodestream.schema import GraphObjectType
+from pandas import Timedelta, Timestamp
+
 
 from .query import Query, QueryBatch
 
@@ -231,7 +232,7 @@ def generate_batch_update_relationship_query_batch(
         return QueryBatch(query, params)
 
     def generate_ttl_match_query(self, config: TimeToLiveConfiguration) -> Query:
-        earliest_allowed_time = datetime.utcnow() - timedelta(
+        earliest_allowed_time = Timestamp.utcnow() - Timedelta(
             hours=config.expiry_in_hours
         )
         params = {"earliest_allowed_time": earliest_allowed_time}

tests/unit/test_ingest_query_builder.py

@@ -1,7 +1,6 @@
-from datetime import datetime
+from unittest.mock import patch
 
 import pytest
-from freezegun import freeze_time
 from hamcrest import assert_that, equal_to, equal_to_ignoring_whitespace
 from nodestream.databases.query_executor import (
     OperationOnNodeIdentity,
@@ -23,14 +22,15 @@
     Neo4jIngestQueryBuilder,
 )
 from nodestream_plugin_neo4j.query import COMMIT_QUERY, Query, QueryBatch
+from pandas import Timestamp
 
 
 @pytest.fixture
 def query_builder():
     return Neo4jIngestQueryBuilder(True)
 
 
-GREATEST_DAY = datetime(1998, 3, 25, 2, 0, 1)
+GREATEST_DAY = Timestamp(1998, 3, 25, 2, 0, 1)
 
 BASIC_NODE_TTL = TimeToLiveConfiguration(
     graph_object_type=GraphObjectType.NODE,
@@ -90,7 +90,7 @@ def query_builder():
     },
 )
 
-
+@patch("pandas.Timestamp.utcnow")
 @pytest.mark.parametrize(
     "ttl,expected_query",
     [
@@ -100,8 +100,8 @@ def query_builder():
         (REL_TTL_WITH_CUSTOM_QUERY, REL_TTL_WITH_CUSTOM_QUERY_EXPECTED_QUERY),
     ],
 )
-@freeze_time("1998-03-25 12:00:01")
-def test_generates_expected_queries(query_builder, ttl, expected_query):
+def test_generates_expected_queries(mocked_utcnow, query_builder, ttl, expected_query):
+    mocked_utcnow.return_value = Timestamp(1998, 3, 25, 12, 0, 1)
     resultant_query = query_builder.generate_ttl_query_from_configuration(ttl)
     assert_that(resultant_query, equal_to(expected_query))