diff --git a/airflow/providers/neo4j/hooks/neo4j.py b/airflow/providers/neo4j/hooks/neo4j.py index 5c888fd3415da..f7a14d0bf7f76 100644 --- a/airflow/providers/neo4j/hooks/neo4j.py +++ b/airflow/providers/neo4j/hooks/neo4j.py @@ -52,15 +52,15 @@ def get_conn(self) -> Neo4jDriver: Function that initiates a new Neo4j connection with username, password and database schema. """ + if self.client is not None: + return self.client + self.connection = self.get_connection(self.neo4j_conn_id) self.extras = self.connection.extra_dejson.copy() self.uri = self.get_uri(self.connection) self.log.info('URI: %s', self.uri) - if self.client is not None: - return self.client - is_encrypted = self.connection.extra_dejson.get('encrypted', False) self.client = GraphDatabase.driver( @@ -76,6 +76,7 @@ def get_uri(self, conn: Connection) -> str: - neo4j_scheme - neo4j:// - certs_self_signed - neo4j+ssc:// - certs_trusted_ca - neo4j+s:// + :param conn: connection object. :return: uri """ @@ -106,7 +107,6 @@ def run(self, query) -> Result: Function to create a neo4j session and execute the query in the session. - :param query: Neo4j query :return: Result """ @@ -114,7 +114,8 @@ def run(self, query) -> Result: if not self.connection.schema: with driver.session() as session: result = session.run(query) + return result.data() else: with driver.session(database=self.connection.schema) as session: result = session.run(query) - return result + return result.data() diff --git a/tests/providers/neo4j/hooks/test_neo4j.py b/tests/providers/neo4j/hooks/test_neo4j.py index 87131c4b5cbc9..44ba7883c1c56 100644 --- a/tests/providers/neo4j/hooks/test_neo4j.py +++ b/tests/providers/neo4j/hooks/test_neo4j.py @@ -15,51 +15,94 @@ # specific language governing permissions and limitations # under the License. # -import json import unittest from unittest import mock +from parameterized import parameterized + from airflow.models import Connection from airflow.providers.neo4j.hooks.neo4j import Neo4jHook class TestNeo4jHookConn(unittest.TestCase): - def setUp(self): - super().setUp() - self.neo4j_hook = Neo4jHook() - self.connection = Connection( - conn_type='neo4j', login='login', password='password', host='host', schema='schema' + @parameterized.expand( + [ + [{}, "bolt://host:7687"], + [{"bolt_scheme": True}, "bolt://host:7687"], + [{"certs_self_signed": True, "bolt_scheme": True}, "bolt+ssc://host:7687"], + [{"certs_trusted_ca": True, "bolt_scheme": True}, "bolt+s://host:7687"], + ] + ) + def test_get_uri_neo4j_scheme(self, conn_extra, expected_uri): + connection = Connection( + conn_type='neo4j', + login='login', + password='password', + host='host', + schema='schema', + extra=conn_extra, ) - def test_get_uri_neo4j_scheme(self): - - self.neo4j_hook.get_connection = mock.Mock() - self.neo4j_hook.get_connection.return_value = self.connection - uri = self.neo4j_hook.get_uri(self.connection) - - assert uri == "bolt://host:7687" - - def test_get_uri_bolt_scheme(self): + # Use the environment variable mocking to test saving the configuration as a URI and + # to avoid mocking Airflow models class + with mock.patch.dict('os.environ', AIRFLOW_CONN_NEO4J_DEFAULT=connection.get_uri()): + neo4j_hook = Neo4jHook() + uri = neo4j_hook.get_uri(connection) - self.connection.extra = json.dumps({"bolt_scheme": True}) - self.neo4j_hook.get_connection = mock.Mock() - self.neo4j_hook.get_connection.return_value = self.connection - uri = self.neo4j_hook.get_uri(self.connection) + assert uri == expected_uri - assert uri == "bolt://host:7687" - - def test_get_uri_bolt_ssc_scheme(self): - self.connection.extra = json.dumps({"certs_self_signed": True, "bolt_scheme": True}) - self.neo4j_hook.get_connection = mock.Mock() - self.neo4j_hook.get_connection.return_value = self.connection - uri = self.neo4j_hook.get_uri(self.connection) + @mock.patch('airflow.providers.neo4j.hooks.neo4j.GraphDatabase') + def test_run_with_schema(self, mock_graph_database): + connection = Connection( + conn_type='neo4j', login='login', password='password', host='host', schema='schema' + ) + mock_sql = mock.MagicMock(name="sql") - assert uri == "bolt+ssc://host:7687" + # Use the environment variable mocking to test saving the configuration as a URI and + # to avoid mocking Airflow models class + with mock.patch.dict('os.environ', AIRFLOW_CONN_NEO4J_DEFAULT=connection.get_uri()): + neo4j_hook = Neo4jHook() + op_result = neo4j_hook.run(mock_sql) + mock_graph_database.assert_has_calls( + [ + mock.call.driver('bolt://host:7687', auth=('login', 'password'), encrypted=False), + mock.call.driver().session(database='schema'), + mock.call.driver().session().__enter__(), + mock.call.driver().session().__enter__().run(mock_sql), + mock.call.driver().session().__enter__().run().data(), + mock.call.driver().session().__exit__(None, None, None), + ] + ) + session = mock_graph_database.driver.return_value.session.return_value.__enter__.return_value + self.assertEqual( + session.run.return_value.data.return_value, + op_result, + ) - def test_get_uri_bolt_trusted_ca_scheme(self): - self.connection.extra = json.dumps({"certs_trusted_ca": True, "bolt_scheme": True}) - self.neo4j_hook.get_connection = mock.Mock() - self.neo4j_hook.get_connection.return_value = self.connection - uri = self.neo4j_hook.get_uri(self.connection) + @mock.patch('airflow.providers.neo4j.hooks.neo4j.GraphDatabase') + def test_run_without_schema(self, mock_graph_database): + connection = Connection( + conn_type='neo4j', login='login', password='password', host='host', schema=None + ) + mock_sql = mock.MagicMock(name="sql") - assert uri == "bolt+s://host:7687" + # Use the environment variable mocking to test saving the configuration as a URI and + # to avoid mocking Airflow models class + with mock.patch.dict('os.environ', AIRFLOW_CONN_NEO4J_DEFAULT=connection.get_uri()): + neo4j_hook = Neo4jHook() + op_result = neo4j_hook.run(mock_sql) + mock_graph_database.assert_has_calls( + [ + mock.call.driver('bolt://host:7687', auth=('login', 'password'), encrypted=False), + mock.call.driver().session(), + mock.call.driver().session().__enter__(), + mock.call.driver().session().__enter__().run(mock_sql), + mock.call.driver().session().__enter__().run().data(), + mock.call.driver().session().__exit__(None, None, None), + ] + ) + session = mock_graph_database.driver.return_value.session.return_value.__enter__.return_value + self.assertEqual( + session.run.return_value.data.return_value, + op_result, + ) diff --git a/tests/providers/neo4j/operators/test_neo4j.py b/tests/providers/neo4j/operators/test_neo4j.py index 39c8d697231f9..320919b5ef2f4 100644 --- a/tests/providers/neo4j/operators/test_neo4j.py +++ b/tests/providers/neo4j/operators/test_neo4j.py @@ -15,27 +15,9 @@ # specific language governing permissions and limitations # under the License. -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. import unittest from unittest import mock -from airflow.models.dag import DAG from airflow.providers.neo4j.operators.neo4j import Neo4jOperator from airflow.utils import timezone @@ -46,16 +28,13 @@ class TestNeo4jOperator(unittest.TestCase): - def setUp(self): - args = {'owner': 'airflow', 'start_date': DEFAULT_DATE} - dag = DAG(TEST_DAG_ID, default_args=args) - self.dag = dag - - @mock.patch('airflow.providers.neo4j.operators.neo4j.Neo4jOperator.get_hook') + @mock.patch('airflow.providers.neo4j.operators.neo4j.Neo4jHook') def test_neo4j_operator_test(self, mock_hook): sql = """ MATCH (tom {name: "Tom Hanks"}) RETURN tom """ - op = Neo4jOperator(task_id='basic_neo4j', sql=sql, dag=self.dag) - op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) + op = Neo4jOperator(task_id='basic_neo4j', sql=sql) + op.execute(mock.MagicMock()) + mock_hook.assert_called_once_with(conn_id='neo4j_default') + mock_hook.return_value.run.assert_called_once_with(sql)