content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import logging
  import os
  import unittest
  import uuid
  
  from mock import MagicMock, patch
  from neo4j import GraphDatabase
  from pyhocon import ConfigFactory
  
  from databuilder.publisher import neo4j_csv_publisher
  from databuilder.publisher.neo4j_csv_publisher import Neo4jCsvPublisher
  
  here = os.path.dirname(__file__)
  
  
  class TestPublish(unittest.TestCase):
  
      def setUp(self) -> None:
          logging.basicConfig(level=logging.INFO)
          self._resource_path = os.path.join(here, '../resources/csv_publisher')
  
      def test_publisher(self) -> None:
          with patch.object(GraphDatabase, 'driver') as mock_driver:
              mock_session = MagicMock()
              mock_driver.return_value.session.return_value = mock_session
  
              mock_transaction = MagicMock()
              mock_session.begin_transaction.return_value = mock_transaction
  
              mock_run = MagicMock()
              mock_transaction.run = mock_run
              mock_commit = MagicMock()
              mock_transaction.commit = mock_commit
  
              publisher = Neo4jCsvPublisher()
  
              conf = ConfigFactory.from_dict(
                  {neo4j_csv_publisher.NEO4J_END_POINT_KEY: 'bolt://999.999.999.999:7687',
                   neo4j_csv_publisher.NODE_FILES_DIR: f'{self._resource_path}/nodes',
                   neo4j_csv_publisher.RELATION_FILES_DIR: f'{self._resource_path}/relations',
                   neo4j_csv_publisher.NEO4J_USER: 'neo4j_user',
                   neo4j_csv_publisher.NEO4J_PASSWORD: 'neo4j_password',
                   neo4j_csv_publisher.JOB_PUBLISH_TAG: str(uuid.uuid4())}
              )
              publisher.init(conf)
              publisher.publish()
  
              self.assertEqual(mock_run.call_count, 6)
  
              # 2 node files, 1 relation file
              self.assertEqual(mock_commit.call_count, 1)
  
      def test_preprocessor(self) -> None:
          with patch.object(GraphDatabase, 'driver') as mock_driver:
              mock_session = MagicMock()
              mock_driver.return_value.session.return_value = mock_session
  
              mock_transaction = MagicMock()
              mock_session.begin_transaction.return_value = mock_transaction
  
              mock_run = MagicMock()
              mock_transaction.run = mock_run
              mock_commit = MagicMock()
              mock_transaction.commit = mock_commit
  
              mock_preprocessor = MagicMock()
              mock_preprocessor.is_perform_preprocess.return_value = MagicMock(return_value=True)
              mock_preprocessor.preprocess_cypher.return_value = ('MATCH (f:Foo) RETURN f', {})
  
              publisher = Neo4jCsvPublisher()
  
              conf = ConfigFactory.from_dict(
                  {neo4j_csv_publisher.NEO4J_END_POINT_KEY: 'bolt://999.999.999.999:7687',
                   neo4j_csv_publisher.NODE_FILES_DIR: f'{self._resource_path}/nodes',
                   neo4j_csv_publisher.RELATION_FILES_DIR: f'{self._resource_path}/relations',
                   neo4j_csv_publisher.RELATION_PREPROCESSOR: mock_preprocessor,
                   neo4j_csv_publisher.NEO4J_USER: 'neo4j_user',
                   neo4j_csv_publisher.NEO4J_PASSWORD: 'neo4j_password',
                   neo4j_csv_publisher.JOB_PUBLISH_TAG: str(uuid.uuid4())}
              )
              publisher.init(conf)
              publisher.publish()
  
              self.assertEqual(mock_run.call_count, 8)
  
              # 2 node files, 1 relation file
              self.assertEqual(mock_commit.call_count, 1)
  
  
  if __name__ == '__main__':
      unittest.main()
  