content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  # Validation of Cypher statements causing Flake8 to fail. Disabling it on this file only
  # flake8: noqa
  
  import logging
  import textwrap
  import unittest
  
  from mock import patch
  from neo4j import GraphDatabase
  from pyhocon import ConfigFactory
  
  from databuilder.publisher import neo4j_csv_publisher
  from databuilder.task import neo4j_staleness_removal_task
  from databuilder.task.neo4j_staleness_removal_task import Neo4jStalenessRemovalTask, TargetWithCondition
  
  
  class TestRemoveStaleData(unittest.TestCase):
  
      def setUp(self) -> None:
          logging.basicConfig(level=logging.INFO)
  
      def test_validation_failure(self) -> None:
  
          with patch.object(GraphDatabase, 'driver'):
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 90,
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo'
              })
  
              task.init(job_config)
              total_record_count = 100
              stale_record_count = 50
              target_type = 'foo'
              task._validate_staleness_pct(total_record_count=total_record_count,
                                           stale_record_count=stale_record_count,
                                           target_type=target_type)
  
      def test_validation(self) -> None:
  
          with patch.object(GraphDatabase, 'driver'):
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo'
              })
  
              task.init(job_config)
              total_record_count = 100
              stale_record_count = 50
              target_type = 'foo'
              self.assertRaises(Exception, task._validate_staleness_pct, total_record_count, stale_record_count, target_type)
  
      def test_validation_threshold_override(self) -> None:
  
          with patch.object(GraphDatabase, 'driver'):
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_PCT_MAX_DICT}': {'foo': 51},
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo'
              })
  
              task.init(job_config)
              task._validate_staleness_pct(total_record_count=100,
                                           stale_record_count=50,
                                           target_type='foo')
              task._validate_staleness_pct(total_record_count=100,
                                           stale_record_count=3,
                                           target_type='bar')
  
      def test_marker(self) -> None:
          with patch.object(GraphDatabase, 'driver'):
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo'
              })
  
              task.init(job_config)
              self.assertIsNone(task.ms_to_expire)
              self.assertEqual(task.marker, 'foo')
  
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 86400000,
              })
  
              task.init(job_config)
              self.assertIsNotNone(task.ms_to_expire)
              self.assertEqual(task.marker, 86400000)
  
      def test_validation_statement_publish_tag(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
              })
  
              task.init(job_config)
              task._validate_node_staleness_pct()
  
              mock_execute.assert_called()
              mock_execute.assert_any_call(statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE true
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(param_dict={'marker': u'foo'},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag))
              RETURN count(*) as count
              """))
  
              task._validate_relation_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': u'foo'},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag))
              RETURN count(*) as count
              """))
  
      def test_validation_statement_publish_tag_retain_data_with_no_publisher_metadata(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.RETAIN_DATA_WITH_NO_PUBLISHER_METADATA}': True
              })
  
              task.init(job_config)
              task._validate_node_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': u'foo'},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.published_tag < $marker)
              RETURN count(*) as count
              """))
  
              task._validate_relation_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': u'foo'},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.published_tag < $marker)
              RETURN count(*) as count
              """))
  
      def test_validation_statement_ms_to_expire(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 9876543210
              })
  
              task.init(job_config)
              task._validate_node_staleness_pct()
  
              mock_execute.assert_called()
              mock_execute.assert_any_call(statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE true
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(param_dict={'marker': 9876543210},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker)
              OR NOT EXISTS(target.publisher_last_updated_epoch_ms))
              RETURN count(*) as count
              """))
  
              task._validate_relation_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': 9876543210},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker)
              OR NOT EXISTS(target.publisher_last_updated_epoch_ms))
              RETURN count(*) as count
              """))
  
      def test_validation_statement_ms_to_expire_retain_data_with_no_publisher_metadata(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 9876543210,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.RETAIN_DATA_WITH_NO_PUBLISHER_METADATA}': True
              })
  
              task.init(job_config)
              task._validate_node_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': 9876543210},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker))
              RETURN count(*) as count
              """))
  
              task._validate_relation_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': 9876543210},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker))
              RETURN count(*) as count
              """))
  
      def test_validation_statement_with_target_condition(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': [TargetWithCondition('Foo', '(target)-[:BAR]->(:Foo) AND target.name=\'foo_name\'')],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': [TargetWithCondition('BAR', '(start_node:Foo)-[target]->(end_node:Foo)')],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
              })
  
              task.init(job_config)
              task._validate_node_staleness_pct()
  
              mock_execute.assert_called()
              mock_execute.assert_any_call(statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE true AND (target)-[:BAR]->(:Foo) AND target.name=\'foo_name\'
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(param_dict={'marker': u'foo'},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag)) AND (target)-[:BAR]->(:Foo) AND target.name=\'foo_name\'
              RETURN count(*) as count
              """))
  
              task._validate_relation_staleness_pct()
              mock_execute.assert_any_call(param_dict={'marker': u'foo'},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag)) AND (start_node:Foo)-[target]->(end_node:Foo)
              RETURN count(*) as count
              """))
  
      def test_validation_receives_correct_counts(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
              })
  
              task.init(job_config)
  
              with patch.object(Neo4jStalenessRemovalTask, '_validate_staleness_pct') as mock_validate:
                  mock_execute.side_effect = [[{'count': 100}], [{'count': 50}]]
                  task._validate_node_staleness_pct()
                  mock_validate.assert_called_with(total_record_count=100,
                                                   stale_record_count=50,
                                                   target_type='Foo')
  
                  mock_execute.side_effect = [[{'count': 100}], [{'count': 50}]]
                  task._validate_relation_staleness_pct()
                  mock_validate.assert_called_with(total_record_count=100,
                                                   stale_record_count=50,
                                                   target_type='BAR')
  
      def test_delete_statement_publish_tag(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              mock_execute.return_value.single.return_value = {'count': 0}
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
              })
  
              task.init(job_config)
              task._delete_stale_nodes()
              task._delete_stale_relations()
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': u'foo', 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag))
              WITH target LIMIT $batch_size
              DETACH DELETE (target)
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': u'foo', 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag))
              WITH target LIMIT $batch_size
              DELETE target
              RETURN count(*) as count
              """))
  
      def test_delete_statement_publish_tag_retain_data_with_no_publisher_metadata(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              mock_execute.return_value.single.return_value = {'count': 0}
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.RETAIN_DATA_WITH_NO_PUBLISHER_METADATA}': True
              })
  
              task.init(job_config)
              task._delete_stale_nodes()
              task._delete_stale_relations()
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': u'foo', 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.published_tag < $marker)
              WITH target LIMIT $batch_size
              DETACH DELETE (target)
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': u'foo', 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.published_tag < $marker)
              WITH target LIMIT $batch_size
              DELETE target
              RETURN count(*) as count
              """))
  
      def test_delete_statement_ms_to_expire(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              mock_execute.return_value.single.return_value = {'count': 0}
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 9876543210
              })
  
              task.init(job_config)
              task._delete_stale_nodes()
              task._delete_stale_relations()
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': 9876543210, 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker)
              OR NOT EXISTS(target.publisher_last_updated_epoch_ms))
              WITH target LIMIT $batch_size
              DETACH DELETE (target)
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': 9876543210, 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker)
              OR NOT EXISTS(target.publisher_last_updated_epoch_ms))
              WITH target LIMIT $batch_size
              DELETE target
              RETURN count(*) as count
              """))
  
      def test_delete_statement_ms_to_expire_retain_data_with_no_publisher_metadata(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              mock_execute.return_value.single.return_value = {'count': 0}
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 9876543210,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.RETAIN_DATA_WITH_NO_PUBLISHER_METADATA}': True
              })
  
              task.init(job_config)
              task._delete_stale_nodes()
              task._delete_stale_relations()
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': 9876543210, 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker))
              WITH target LIMIT $batch_size
              DETACH DELETE (target)
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': 9876543210, 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.publisher_last_updated_epoch_ms < (timestamp() - $marker))
              WITH target LIMIT $batch_size
              DELETE target
              RETURN count(*) as count
              """))
  
      def test_delete_statement_with_target_condition(self) -> None:
          with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jStalenessRemovalTask, '_execute_cypher_query') \
                  as mock_execute:
              mock_execute.return_value.single.return_value = {'count': 0}
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': [TargetWithCondition('Foo', '(target)-[:BAR]->(:Foo) AND target.name=\'foo_name\'')],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': [TargetWithCondition('BAR', '(start_node:Foo)-[target]->(end_node:Foo)')],
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
              })
  
              task.init(job_config)
              task._delete_stale_nodes()
              task._delete_stale_relations()
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': u'foo', 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (target:Foo)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag)) AND (target)-[:BAR]->(:Foo) AND target.name=\'foo_name\'
              WITH target LIMIT $batch_size
              DETACH DELETE (target)
              RETURN count(*) as count
              """))
  
              mock_execute.assert_any_call(dry_run=False,
                                           param_dict={'marker': u'foo', 'batch_size': 100},
                                           statement=textwrap.dedent("""
              MATCH (start_node)-[target:BAR]-(end_node)
              WHERE (target.published_tag < $marker
              OR NOT EXISTS(target.published_tag)) AND (start_node:Foo)-[target]->(end_node:Foo)
              WITH target LIMIT $batch_size
              DELETE target
              RETURN count(*) as count
              """))
  
      def test_ms_to_expire_too_small(self) -> None:
          with patch.object(GraphDatabase, 'driver'):
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 24 * 60 * 60 * 100 - 10
              })
  
              try:
                  task.init(job_config)
                  self.assertTrue(False, 'Should have failed with small TTL   ')
              except Exception:
                  pass
  
          with patch.object(GraphDatabase, 'driver'):
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.MS_TO_EXPIRE}': 24 * 60 * 60 * 1000,
              })
              task.init(job_config)
  
      def test_delete_dry_run(self) -> None:
          with patch.object(GraphDatabase, 'driver') as mock_driver:
              session_mock = mock_driver.return_value.session
  
              task = Neo4jStalenessRemovalTask()
              job_config = ConfigFactory.from_dict({
                  f'job.identifier': 'remove_stale_data_job',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_END_POINT_KEY}': 'neo4j://example.com:7687',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_USER}': 'foo',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.NEO4J_PASSWORD}': 'bar',
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.STALENESS_MAX_PCT}': 5,
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_NODES}': ['Foo'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.TARGET_RELATIONS}': ['BAR'],
                  f'{task.get_scope()}.{neo4j_staleness_removal_task.DRY_RUN}': True,
                  neo4j_csv_publisher.JOB_PUBLISH_TAG: 'foo',
              })
  
  