diff --git a/nutkit/protocol/feature.py b/nutkit/protocol/feature.py index c0381584..7a162dd6 100644 --- a/nutkit/protocol/feature.py +++ b/nutkit/protocol/feature.py @@ -220,6 +220,10 @@ class Feature(Enum): CONF_HINT_CON_RECV_TIMEOUT = "ConfHint:connection.recv_timeout_seconds" # === BACKEND FEATURES FOR TESTING === + # The backend/driver offers a way to configure a driver with a custom DNS + # resolver. This configuration option is for testing purposes and might + # not be exposed to the user. + BACKEND_DNS_RESOLVER = "Backend:DNSResolver" # The backend understands the FakeTimeInstall, FakeTimeUninstall and # FakeTimeTick protocol messages and provides a way to mock the system # time. This is mainly used for testing various timeouts. diff --git a/tests/stub/advertised_address/__init__.py b/tests/stub/advertised_address/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/stub/advertised_address/scripts/advertised_address_direct.script b/tests/stub/advertised_address/scripts/advertised_address_direct.script new file mode 100644 index 00000000..c02c2f48 --- /dev/null +++ b/tests/stub/advertised_address/scripts/advertised_address_direct.script @@ -0,0 +1,24 @@ +!: BOLT 5.8 + +A: HELLO {"{}": "*"} +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"} + +*: RESET + +C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [1] +S: SUCCESS {} + +*: RESET + +C: RUN "RETURN 2 AS n" {"{}": "*"} {"{}": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [2] +S: SUCCESS {} + +*: RESET +A: GOODBYE diff --git a/tests/stub/advertised_address/scripts/advertised_address_direct_warm.script b/tests/stub/advertised_address/scripts/advertised_address_direct_warm.script new file mode 100644 index 00000000..3a41ae4b --- /dev/null +++ b/tests/stub/advertised_address/scripts/advertised_address_direct_warm.script @@ -0,0 +1,44 @@ +!: BOLT 5.8 + +A: HELLO {"{}": "*"} +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "something_wild:1234"} + +*: RESET + +C: RUN "RETURN 1 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [1] +S: SUCCESS {} + +*: RESET + +A: LOGOFF +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "#HOST#:#PORT#"} + +*: RESET + +C: RUN "RETURN 2 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [2] +S: SUCCESS {} + +*: RESET + +A: LOGOFF +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "certainly-unusable:7688"} + +*: RESET + +C: RUN "RETURN 3 AS n" {"{}": "*"} {"mode": "r", "[db]": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [3] +S: SUCCESS {} + +*: RESET +A: GOODBYE diff --git a/tests/stub/advertised_address/scripts/advertised_address_routing.script b/tests/stub/advertised_address/scripts/advertised_address_routing.script new file mode 100644 index 00000000..18a28011 --- /dev/null +++ b/tests/stub/advertised_address/scripts/advertised_address_routing.script @@ -0,0 +1,21 @@ +!: BOLT 5.8 + +A: HELLO {"{}": "*"} +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"} + +*: RESET + +C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"} +S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["doesnt_exist2:#PORT#"], "role":"READ"}, {"addresses": ["#ADVERTISED_HOST#:#PORT#"], "role":"WRITE"}]}} + +*: RESET + +C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [1] +S: SUCCESS {} + +*: RESET +A: GOODBYE diff --git a/tests/stub/advertised_address/scripts/advertised_address_routing_warm.script b/tests/stub/advertised_address/scripts/advertised_address_routing_warm.script new file mode 100644 index 00000000..da8c5cdb --- /dev/null +++ b/tests/stub/advertised_address/scripts/advertised_address_routing_warm.script @@ -0,0 +1,49 @@ +!: BOLT 5.8 + +A: HELLO {"{}": "*"} +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "#HOST#:#PORT#"} + +*: RESET + +C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"} +S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["#ADVERTISED_HOST#:#PORT#"], "role":"READ"}, {"addresses": ["#HOST#:#PORT#"], "role":"WRITE"}]}} + +*: RESET + +C: RUN "RETURN 1 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [1] +S: SUCCESS {} + +*: RESET + +A: LOGOFF +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"} + +*: RESET + +C: RUN "RETURN 2 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [2] +S: SUCCESS {} + +*: RESET + +A: LOGOFF +C: LOGON {"{}": "*"} +S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"} + +*: RESET + +C: RUN "RETURN 3 AS n" {"{}": "*"} {"mode": "r", "[db]": "*"} +S: SUCCESS {"fields": ["n"]} +C: PULL {"{}": "*"} +S: RECORD [3] +S: SUCCESS {} + +*: RESET +A: GOODBYE diff --git a/tests/stub/advertised_address/test_advertised_address.py b/tests/stub/advertised_address/test_advertised_address.py new file mode 100644 index 00000000..6caa22da --- /dev/null +++ b/tests/stub/advertised_address/test_advertised_address.py @@ -0,0 +1,198 @@ +from abc import ABC +from collections import deque +from contextlib import contextmanager + +import nutkit.protocol as types +from nutkit.frontend import Driver +from tests.shared import ( + driver_feature, + TestkitTestCase, +) +from tests.stub.shared import StubServer + +_FAKE_ADDRESS = "banana.example.com" +_FAKE_ADVERTISED_ADDRESS = "cucumber.example.com" + + +class _AdvertisedAddressTestCase(TestkitTestCase, ABC): + @contextmanager + def server(self, script, vars_=None): + server = StubServer(9001) + vars_ = { + "#ADVERTISED_HOST#": f"{_FAKE_ADVERTISED_ADDRESS}", + "#PORT#": server.port, + "#HOST#": server.host, + **(vars_ or {}), + } + server.start(path=self.script_path(script), + vars_=vars_) + try: + yield server + except Exception: + server.reset() + raise + + server.done() + + @contextmanager + def driver(self, server, routing=True, dns_resolver=None, **kwargs): + auth = types.AuthorizationToken("bearer", credentials="foo") + scheme = "neo4j" if routing else "bolt" + uri = f"{scheme}://{_FAKE_ADDRESS}:{server.port}" + driver = Driver( + self._backend, uri, auth, domain_name_resolver_fn=dns_resolver, + **kwargs, + ) + try: + yield driver + finally: + driver.close() + + @contextmanager + def session(self, driver, access_mode="w", session_config=None): + if session_config is None: + session_config = {} + session = driver.session(access_mode, **session_config) + try: + yield session + finally: + session.close() + + @staticmethod + def _make_dns_resolver(*expected_resolved_pairs): + dns_expectations = deque(expected_resolved_pairs) + + def dns_resolver(name): + nonlocal dns_expectations + expectation, result = dns_expectations.popleft() + parts = name.rsplit(":", 1) + sep = port = "" + if len(parts) == 2: + name, port = parts + sep = ":" + print(parts, expectation) + assert name == expectation + return [sep.join((host, port)) for host in result] + + return dns_resolver + + +class TestAdvertisedAddress(_AdvertisedAddressTestCase): + required_features = ( + types.Feature.BOLT_5_8, + ) + + def test_routing_driver_reuses_connection_according_to_advertised_address( + self, + ): + with self.server("advertised_address_routing.script") as server: + self._test_reuses_connection( + server, + driver_kwargs={"routing": True}, + ) + + def test_direct_driver_reuses_connection_regardless_of_advertised_address( + self + ): + with self.server("advertised_address_direct.script") as server: + self._test_reuses_connection( + server, + driver_kwargs={ + "routing": False, + "max_connection_pool_size": 1, + }, + repetitions=2, + ) + + @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) + def _test_reuses_connection( + self, + server, + *, + driver_kwargs=None, + repetitions=1, + ): + if driver_kwargs is None: + driver_kwargs = {} + + dns_resolver = self._make_dns_resolver((_FAKE_ADDRESS, [server.host])) + + with self.driver( + server, + dns_resolver=dns_resolver, + **driver_kwargs + ) as driver: + for i in range(repetitions): + with self.session(driver) as session: + list(session.run(f"RETURN {i + 1} AS n")) + + @driver_feature( + types.Feature.BACKEND_DNS_RESOLVER, + types.Feature.API_SESSION_AUTH_CONFIG, + ) + def test_warm_routing_driver_reuses_connection_according_to_advertised_address( # noqa: E501 + self, + ): + with self.server("advertised_address_routing_warm.script") as server: + self._test_reuses_warm_connection( + server, + driver_kwargs={ + "routing": True, + "max_connection_pool_size": 1, + } + ) + + @driver_feature( + types.Feature.BACKEND_DNS_RESOLVER, + types.Feature.API_SESSION_AUTH_CONFIG, + ) + def test_warm_direct_driver_reuses_connection_according_to_advertised_address( # noqa: E501 + self, + ): + with self.server("advertised_address_direct_warm.script") as server: + self._test_reuses_warm_connection( + server, + driver_kwargs={ + "routing": False, + "max_connection_pool_size": 1, + } + ) + + @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) + def _test_reuses_warm_connection( + self, + server, + *, + driver_kwargs=None, + ): + if driver_kwargs is None: + driver_kwargs = {} + + dns_resolver = self._make_dns_resolver((_FAKE_ADDRESS, [server.host])) + + with self.driver( + server, + dns_resolver=dns_resolver, + **driver_kwargs + ) as driver: + with self.session( + driver, + session_config={"database": "neo4j"}, + ) as session: + list(session.run("RETURN 1 AS n")) + + # Using session auth to cause LOGOFF/LOGON + # during which the server will change its advertised address. + auth = types.AuthorizationToken("bearer", credentials="bar") + with self.session( + driver, + session_config={"database": "neo4j", "auth_token": auth} + ) as session: + list(session.run("RETURN 2 AS n")) + + with self.session( + driver, + access_mode="r", + session_config={"database": "neo4j"}, + ) as session: + list(session.run("RETURN 3 AS n")) diff --git a/tests/stub/routing/test_routing_v5x0.py b/tests/stub/routing/test_routing_v5x0.py index b628a7e5..00d0ef2b 100644 --- a/tests/stub/routing/test_routing_v5x0.py +++ b/tests/stub/routing/test_routing_v5x0.py @@ -2412,6 +2412,7 @@ def test_should_ignore_system_bookmark_when_getting_rt_for_multi_db(self): self.assertEqual([1], sequence) self.assertEqual(["foo:6678"], last_bookmarks) + @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) def _test_should_request_rt_from_all_initial_routers_until_successful( self, failure_script ):