From e0e6cf2d8385cc87da866c30652eb87c07295394 Mon Sep 17 00:00:00 2001 From: Stephen Cathcart Date: Mon, 9 Dec 2024 10:42:13 +0000 Subject: [PATCH 1/5] Add tests for advertised address Co-authored-by: Robsdedude --- boltstub/bolt_protocol.py | 9 +++ nutkit/protocol/feature.py | 2 + tests/stub/advertised_address/__init__.py | 0 .../scripts/advertised_address.script | 19 +++++ .../test_advertised_address.py | 81 +++++++++++++++++++ 5 files changed, 111 insertions(+) create mode 100644 tests/stub/advertised_address/__init__.py create mode 100644 tests/stub/advertised_address/scripts/advertised_address.script create mode 100644 tests/stub/advertised_address/test_advertised_address.py diff --git a/boltstub/bolt_protocol.py b/boltstub/bolt_protocol.py index fb70c08d2..61fc012e8 100644 --- a/boltstub/bolt_protocol.py +++ b/boltstub/bolt_protocol.py @@ -570,3 +570,12 @@ class Bolt5x7Protocol(Bolt5x6Protocol): equivalent_versions = set() server_agent = "Neo4j/5.24.0" + + +class Bolt5x8Protocol(Bolt5x7Protocol): + protocol_version = (5, 8) + version_aliases = set() + # allow the server to negotiate other bolt versions + equivalent_versions = set() + + server_agent = "Neo4j/5.26.0" diff --git a/nutkit/protocol/feature.py b/nutkit/protocol/feature.py index 679658a3b..2a526882f 100644 --- a/nutkit/protocol/feature.py +++ b/nutkit/protocol/feature.py @@ -125,6 +125,8 @@ class Feature(Enum): BOLT_5_6 = "Feature:Bolt:5.6" # The driver supports Bolt protocol version 5.7 BOLT_5_7 = "Feature:Bolt:5.7" + # The driver supports Bolt protocol version 5.8 + BOLT_5_8 = "Feature:Bolt:5.8" # The driver supports patching DateTimes to use UTC for Bolt 4.3 and 4.4 BOLT_PATCH_UTC = "Feature:Bolt:Patch:UTC" # The driver supports impersonation diff --git a/tests/stub/advertised_address/__init__.py b/tests/stub/advertised_address/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/stub/advertised_address/scripts/advertised_address.script b/tests/stub/advertised_address/scripts/advertised_address.script new file mode 100644 index 000000000..ba3d4fbe6 --- /dev/null +++ b/tests/stub/advertised_address/scripts/advertised_address.script @@ -0,0 +1,19 @@ +!: BOLT 5.8 + +A: HELLO {"{}": "*"} +C: LOGON {"{}": "*"} +S: SUCCESS {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_ADDRESS#"} + +C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"} +S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["doesnt_exist2:9001"], "role":"READ"}, {"addresses": ["#ADVERTISED_ADDRESS#:9001"], "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/test_advertised_address.py b/tests/stub/advertised_address/test_advertised_address.py new file mode 100644 index 000000000..711383814 --- /dev/null +++ b/tests/stub/advertised_address/test_advertised_address.py @@ -0,0 +1,81 @@ +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 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): + if vars_ is None: + vars_ = {} + server = StubServer(9001) + 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): + 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, + ) + try: + yield driver + finally: + driver.close() + + @contextmanager + def session(self, driver): + session = driver.session("w") + try: + yield session + finally: + session.close() + + +class TestAdvertisedAddress(_AdvertisedAddressTestCase): + required_features = ( + types.Feature.BOLT_5_8, + ) + + def test_advertised_address(self): + vars_ = { + "#ADVERTISED_ADDRESS#": _FAKE_ADVERTISED_ADDRESS, + } + with self.server("advertised_address.script", vars_=vars_) as server: + + dns_expectations = deque( + ( + _FAKE_ADDRESS, + [server.host], + ), + ) + + def dns_resolver(name): + nonlocal dns_expectations + expectation, result = dns_expectations.popleft() + assert name == expectation[0] + return result + + with self.driver(server, dns_resolver=dns_resolver) as driver: + with self.session(driver) as session: + list(session.run("RETURN 1 AS n")) + + # TODO: test direct connection should be re-used regardless + # of advertised address From 882c64c030d6905b0410683e9c646da3ba97fb3a Mon Sep 17 00:00:00 2001 From: Robsdedude Date: Tue, 14 Jan 2025 10:42:49 +0100 Subject: [PATCH 2/5] Fix advertised address test --- .../scripts/advertised_address.script | 4 +-- .../test_advertised_address.py | 25 +++++++++++-------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/tests/stub/advertised_address/scripts/advertised_address.script b/tests/stub/advertised_address/scripts/advertised_address.script index ba3d4fbe6..4037b2616 100644 --- a/tests/stub/advertised_address/scripts/advertised_address.script +++ b/tests/stub/advertised_address/scripts/advertised_address.script @@ -2,10 +2,10 @@ A: HELLO {"{}": "*"} C: LOGON {"{}": "*"} -S: SUCCESS {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_ADDRESS#"} +S: SUCCESS {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_HOST#:#ADVERTISED_PORT#"} C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"} -S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["doesnt_exist2:9001"], "role":"READ"}, {"addresses": ["#ADVERTISED_ADDRESS#:9001"], "role":"WRITE"}]}} +S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["doesnt_exist2:#ADVERTISED_PORT#"], "role":"READ"}, {"addresses": ["#ADVERTISED_HOST#:#ADVERTISED_PORT#"], "role":"WRITE"}]}} *: RESET diff --git a/tests/stub/advertised_address/test_advertised_address.py b/tests/stub/advertised_address/test_advertised_address.py index 711383814..656ab5a83 100644 --- a/tests/stub/advertised_address/test_advertised_address.py +++ b/tests/stub/advertised_address/test_advertised_address.py @@ -14,9 +14,12 @@ class _AdvertisedAddressTestCase(TestkitTestCase, ABC): @contextmanager def server(self, script, vars_=None): - if vars_ is None: - vars_ = {} server = StubServer(9001) + vars_ = { + "#ADVERTISED_HOST#": f"{_FAKE_ADVERTISED_ADDRESS}", + "#ADVERTISED_PORT#": server.port, + **(vars_ or {}), + } server.start(path=self.script_path(script), vars_=vars_) try: @@ -55,23 +58,23 @@ class TestAdvertisedAddress(_AdvertisedAddressTestCase): ) def test_advertised_address(self): - vars_ = { - "#ADVERTISED_ADDRESS#": _FAKE_ADVERTISED_ADDRESS, - } - with self.server("advertised_address.script", vars_=vars_) as server: + with self.server("advertised_address.script") as server: dns_expectations = deque( ( - _FAKE_ADDRESS, - [server.host], - ), + ( + _FAKE_ADDRESS, + [server.host], + ), + ) ) def dns_resolver(name): nonlocal dns_expectations expectation, result = dns_expectations.popleft() - assert name == expectation[0] - return result + name, sep, port = name.rpartition(":") + assert name == expectation + return [sep.join((host, port)) for host in result] with self.driver(server, dns_resolver=dns_resolver) as driver: with self.session(driver) as session: From 4697fd1a58f2f1cc567c47b017a3d8a9f5ea5b69 Mon Sep 17 00:00:00 2001 From: Robsdedude Date: Tue, 14 Jan 2025 11:20:17 +0100 Subject: [PATCH 3/5] Introduce feature flag for DNS resolution hook --- nutkit/protocol/feature.py | 4 ++++ tests/stub/advertised_address/test_advertised_address.py | 6 +++++- tests/stub/routing/test_routing_v5x0.py | 1 + 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/nutkit/protocol/feature.py b/nutkit/protocol/feature.py index c03815846..7a162dd63 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/test_advertised_address.py b/tests/stub/advertised_address/test_advertised_address.py index 656ab5a83..f7f4b7376 100644 --- a/tests/stub/advertised_address/test_advertised_address.py +++ b/tests/stub/advertised_address/test_advertised_address.py @@ -4,7 +4,10 @@ import nutkit.protocol as types from nutkit.frontend import Driver -from tests.shared import TestkitTestCase +from tests.shared import ( + driver_feature, + TestkitTestCase, +) from tests.stub.shared import StubServer _FAKE_ADDRESS = "banana.example.com" @@ -57,6 +60,7 @@ class TestAdvertisedAddress(_AdvertisedAddressTestCase): types.Feature.BOLT_5_8, ) + @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) def test_advertised_address(self): with self.server("advertised_address.script") as server: diff --git a/tests/stub/routing/test_routing_v5x0.py b/tests/stub/routing/test_routing_v5x0.py index b628a7e55..00d0ef2b0 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 ): From 8d6549811389910d1b270001260320c77d62a1af Mon Sep 17 00:00:00 2001 From: Robsdedude Date: Tue, 14 Jan 2025 11:46:01 +0100 Subject: [PATCH 4/5] Add advertised address test for direct driver --- .../scripts/advertised_address_direct.script | 24 +++++++ ...ript => advertised_address_routing.script} | 2 + .../test_advertised_address.py | 68 +++++++++++++------ 3 files changed, 74 insertions(+), 20 deletions(-) create mode 100644 tests/stub/advertised_address/scripts/advertised_address_direct.script rename tests/stub/advertised_address/scripts/{advertised_address.script => advertised_address_routing.script} (98%) 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 000000000..da7baa2e8 --- /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 {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_HOST#:#ADVERTISED_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.script b/tests/stub/advertised_address/scripts/advertised_address_routing.script similarity index 98% rename from tests/stub/advertised_address/scripts/advertised_address.script rename to tests/stub/advertised_address/scripts/advertised_address_routing.script index 4037b2616..6190d7c16 100644 --- a/tests/stub/advertised_address/scripts/advertised_address.script +++ b/tests/stub/advertised_address/scripts/advertised_address_routing.script @@ -4,6 +4,8 @@ A: HELLO {"{}": "*"} C: LOGON {"{}": "*"} S: SUCCESS {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_HOST#:#ADVERTISED_PORT#"} +*: RESET + C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"} S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["doesnt_exist2:#ADVERTISED_PORT#"], "role":"READ"}, {"addresses": ["#ADVERTISED_HOST#:#ADVERTISED_PORT#"], "role":"WRITE"}]}} diff --git a/tests/stub/advertised_address/test_advertised_address.py b/tests/stub/advertised_address/test_advertised_address.py index f7f4b7376..91fda7e01 100644 --- a/tests/stub/advertised_address/test_advertised_address.py +++ b/tests/stub/advertised_address/test_advertised_address.py @@ -34,12 +34,13 @@ def server(self, script, vars_=None): server.done() @contextmanager - def driver(self, server, routing=True, dns_resolver=None): + 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 @@ -60,29 +61,56 @@ class TestAdvertisedAddress(_AdvertisedAddressTestCase): types.Feature.BOLT_5_8, ) - @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) - def test_advertised_address(self): - with self.server("advertised_address.script") as server: + def _test_reuses_connection( + self, + server, + *, + driver_kwargs=None, + repetitions=1, + ): + if driver_kwargs is None: + driver_kwargs = {} - dns_expectations = deque( + dns_expectations = deque( + ( ( - ( - _FAKE_ADDRESS, - [server.host], - ), - ) + _FAKE_ADDRESS, + [server.host], + ), ) + ) - def dns_resolver(name): - nonlocal dns_expectations - expectation, result = dns_expectations.popleft() - name, sep, port = name.rpartition(":") - assert name == expectation - return [sep.join((host, port)) for host in result] + def dns_resolver(name): + nonlocal dns_expectations + expectation, result = dns_expectations.popleft() + name, sep, port = name.rpartition(":") + assert name == expectation + return [sep.join((host, port)) for host in result] - with self.driver(server, dns_resolver=dns_resolver) as driver: + 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("RETURN 1 AS n")) + list(session.run(f"RETURN {i + 1} AS n")) + + @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) + def test_reuses_connection_according_to_advertised_address_routing(self): + with self.server("advertised_address_routing.script") as server: + self._test_reuses_connection( + server, + driver_kwargs={"routing": True}, + ) - # TODO: test direct connection should be re-used regardless - # of advertised address + def test_reuses_connection_regardless_of_advertised_address_direct(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, + ) From 8a9fdb02f894214cac12a5ad234ff28cf975a822 Mon Sep 17 00:00:00 2001 From: Robsdedude Date: Tue, 14 Jan 2025 15:55:22 +0100 Subject: [PATCH 5/5] Add tests for advertised address changing on re-auth --- .../scripts/advertised_address_direct.script | 2 +- .../advertised_address_direct_warm.script | 44 ++++++ .../scripts/advertised_address_routing.script | 4 +- .../advertised_address_routing_warm.script | 49 +++++++ .../test_advertised_address.py | 138 ++++++++++++++---- 5 files changed, 206 insertions(+), 31 deletions(-) create mode 100644 tests/stub/advertised_address/scripts/advertised_address_direct_warm.script create mode 100644 tests/stub/advertised_address/scripts/advertised_address_routing_warm.script diff --git a/tests/stub/advertised_address/scripts/advertised_address_direct.script b/tests/stub/advertised_address/scripts/advertised_address_direct.script index da7baa2e8..c02c2f489 100644 --- a/tests/stub/advertised_address/scripts/advertised_address_direct.script +++ b/tests/stub/advertised_address/scripts/advertised_address_direct.script @@ -2,7 +2,7 @@ A: HELLO {"{}": "*"} C: LOGON {"{}": "*"} -S: SUCCESS {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_HOST#:#ADVERTISED_PORT#"} +S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"} *: RESET 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 000000000..3a41ae4b0 --- /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 index 6190d7c16..18a280110 100644 --- a/tests/stub/advertised_address/scripts/advertised_address_routing.script +++ b/tests/stub/advertised_address/scripts/advertised_address_routing.script @@ -2,12 +2,12 @@ A: HELLO {"{}": "*"} C: LOGON {"{}": "*"} -S: SUCCESS {"connection_id": "bolt-0", "server": "Neo4j/5.26.0", "advertised_address": "#ADVERTISED_HOST#:#ADVERTISED_PORT#"} +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:#ADVERTISED_PORT#"], "role":"READ"}, {"addresses": ["#ADVERTISED_HOST#:#ADVERTISED_PORT#"], "role":"WRITE"}]}} +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 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 000000000..da8c5cdb1 --- /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 index 91fda7e01..6caa22dab 100644 --- a/tests/stub/advertised_address/test_advertised_address.py +++ b/tests/stub/advertised_address/test_advertised_address.py @@ -20,7 +20,8 @@ def server(self, script, vars_=None): server = StubServer(9001) vars_ = { "#ADVERTISED_HOST#": f"{_FAKE_ADVERTISED_ADDRESS}", - "#ADVERTISED_PORT#": server.port, + "#PORT#": server.port, + "#HOST#": server.host, **(vars_ or {}), } server.start(path=self.script_path(script), @@ -48,19 +49,62 @@ def driver(self, server, routing=True, dns_resolver=None, **kwargs): driver.close() @contextmanager - def session(self, driver): - session = driver.session("w") + 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, @@ -71,21 +115,7 @@ def _test_reuses_connection( if driver_kwargs is None: driver_kwargs = {} - dns_expectations = deque( - ( - ( - _FAKE_ADDRESS, - [server.host], - ), - ) - ) - - def dns_resolver(name): - nonlocal dns_expectations - expectation, result = dns_expectations.popleft() - name, sep, port = name.rpartition(":") - assert name == expectation - return [sep.join((host, port)) for host in result] + dns_resolver = self._make_dns_resolver((_FAKE_ADDRESS, [server.host])) with self.driver( server, @@ -96,21 +126,73 @@ def dns_resolver(name): with self.session(driver) as session: list(session.run(f"RETURN {i + 1} AS n")) - @driver_feature(types.Feature.BACKEND_DNS_RESOLVER) - def test_reuses_connection_according_to_advertised_address_routing(self): - with self.server("advertised_address_routing.script") as server: - self._test_reuses_connection( + @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}, + driver_kwargs={ + "routing": True, + "max_connection_pool_size": 1, + } ) - def test_reuses_connection_regardless_of_advertised_address_direct(self): - with self.server("advertised_address_direct.script") as server: - self._test_reuses_connection( + @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, - }, - repetitions=2, + } ) + + @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"))