Skip to content

Commit b1899d6

Browse files
sleepySticksleepyStick
committed
WIP - sync test pass but not async..
1 parent 2759379 commit b1899d6

File tree

6 files changed

+356
-0
lines changed

6 files changed

+356
-0
lines changed

‎pymongo/asynchronous/mongo_client.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
from pymongo.asynchronous.settings import TopologySettings
7171
from pymongo.asynchronous.topology import Topology, _ErrorContext
7272
from pymongo.client_options import ClientOptions
73+
from pymongo.driver_info import DriverInfo
7374
from pymongo.errors import (
7475
AutoReconnect,
7576
BulkWriteError,
@@ -1040,6 +1041,20 @@ async def target() -> bool:
10401041
self._kill_cursors_executor = executor
10411042
self._opened = False
10421043

1044+
def _append_metadata(self, driver_info: DriverInfo) -> None:
1045+
metadata = self._options.pool_options.metadata
1046+
for k, v in driver_info._asdict().items():
1047+
if v is None:
1048+
continue
1049+
if k in metadata:
1050+
metadata[k] = f"{metadata[k]}|{v}"
1051+
elif k in metadata["driver"]:
1052+
metadata["driver"][k] = "{}|{}".format(
1053+
metadata["driver"][k],
1054+
v,
1055+
)
1056+
self._options.pool_options._set_metadata(metadata)
1057+
10431058
def _should_pin_cursor(self, session: Optional[AsyncClientSession]) -> Optional[bool]:
10441059
return self._options.load_balanced and not (session and session.in_transaction)
10451060

‎pymongo/pool_options.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,3 +522,6 @@ def server_api(self) -> Optional[ServerApi]:
522522
def load_balanced(self) -> Optional[bool]:
523523
"""True if this Pool is configured in load balanced mode."""
524524
return self.__load_balanced
525+
526+
def _set_metadata(self, new_data: dict[str, Any]) -> None:
527+
self.__metadata = new_data

‎pymongo/synchronous/mongo_client.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262
from bson.timestamp import Timestamp
6363
from pymongo import _csot, common, helpers_shared, periodic_executor
6464
from pymongo.client_options import ClientOptions
65+
from pymongo.driver_info import DriverInfo
6566
from pymongo.errors import (
6667
AutoReconnect,
6768
BulkWriteError,
@@ -1040,6 +1041,20 @@ def target() -> bool:
10401041
self._kill_cursors_executor = executor
10411042
self._opened = False
10421043

1044+
def _append_metadata(self, driver_info: DriverInfo) -> None:
1045+
metadata = self._options.pool_options.metadata
1046+
for k, v in driver_info._asdict().items():
1047+
if v is None:
1048+
continue
1049+
if k in metadata:
1050+
metadata[k] = f"{metadata[k]}|{v}"
1051+
elif k in metadata["driver"]:
1052+
metadata["driver"][k] = "{}|{}".format(
1053+
metadata["driver"][k],
1054+
v,
1055+
)
1056+
self._options.pool_options._set_metadata(metadata)
1057+
10431058
def _should_pin_cursor(self, session: Optional[ClientSession]) -> Optional[bool]:
10441059
return self._options.load_balanced and not (session and session.in_transaction)
10451060

‎test/asynchronous/test_client_metadata.py

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
# Copyright 2013-present MongoDB, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
from __future__ import annotations
15+
16+
import asyncio
17+
import time
18+
import unittest
19+
from test.utils_shared import CMAPListener
20+
from typing import Optional
21+
22+
import pytest
23+
from mockupdb import MockupDB, OpMsgReply
24+
25+
from pymongo import AsyncMongoClient, MongoClient
26+
from pymongo.driver_info import DriverInfo
27+
from pymongo.monitoring import ConnectionClosedEvent
28+
29+
pytestmark = pytest.mark.mockupdb
30+
31+
_IS_SYNC = False
32+
33+
34+
def _get_handshake_driver_info(request):
35+
assert "client" in request
36+
return request["client"]
37+
38+
39+
class TestClientMetadataProse(unittest.TestCase):
40+
def setUp(self):
41+
self.server = MockupDB()
42+
self.handshake_req = None
43+
44+
def respond(r):
45+
# Only save the very first request from the driver.
46+
if self.handshake_req is None:
47+
self.handshake_req = r
48+
return r.reply(OpMsgReply(minWireVersion=0, maxWireVersion=13))
49+
50+
self.server.autoresponds(respond)
51+
self.server.run()
52+
self.addCleanup(self.server.stop)
53+
54+
async def check_metadata_added(
55+
self, add_name: Optional[str], add_version: Optional[str], add_platform: Optional[str]
56+
) -> None:
57+
client = AsyncMongoClient(
58+
"mongodb://" + self.server.address_string,
59+
maxIdleTimeMS=1,
60+
driver=DriverInfo("library", "1.2", "Library Platform"),
61+
)
62+
63+
# send initial metadata
64+
await client.admin.command("ping")
65+
metadata = _get_handshake_driver_info(self.handshake_req)
66+
driver_metadata = metadata["driver"]
67+
name, version, platform = (
68+
driver_metadata["name"],
69+
driver_metadata["version"],
70+
metadata["platform"],
71+
)
72+
await asyncio.sleep(0.005)
73+
74+
# add data
75+
client._append_metadata(DriverInfo(add_name, add_version, add_platform))
76+
# reset
77+
self.handshake_req = None
78+
await client.admin.command("ping")
79+
new_metadata = _get_handshake_driver_info(self.handshake_req)
80+
# compare
81+
self.assertEqual(
82+
new_metadata["driver"]["name"], f"{name}|{add_name}" if add_name is not None else name
83+
)
84+
self.assertEqual(
85+
new_metadata["driver"]["version"],
86+
f"{version}|{add_version}" if add_version is not None else version,
87+
)
88+
self.assertEqual(
89+
new_metadata["platform"],
90+
f"{platform}|{add_platform}" if add_platform is not None else platform,
91+
)
92+
93+
metadata.pop("driver")
94+
metadata.pop("platform")
95+
new_metadata.pop("driver")
96+
new_metadata.pop("platform")
97+
self.assertEqual(metadata, new_metadata)
98+
99+
await client.close()
100+
101+
async def test_append_metadata(self):
102+
await self.check_metadata_added("framework", "2.0", "Framework Platform")
103+
104+
async def test_append_metadata_platform_none(self):
105+
await self.check_metadata_added("framework", "2.0", None)
106+
107+
async def test_append_metadata_version_none(self):
108+
await self.check_metadata_added("framework", None, "Framework Platform")
109+
110+
async def test_append_platform_version_none(self):
111+
await self.check_metadata_added("framework", None, None)
112+
113+
async def test_doesnt_update_established_connections(self):
114+
listener = CMAPListener()
115+
client = AsyncMongoClient(
116+
"mongodb://" + self.server.address_string,
117+
maxIdleTimeMS=1,
118+
driver=DriverInfo("library", "1.2", "Library Platform"),
119+
event_listeners=[listener],
120+
)
121+
122+
# send initial metadata
123+
await client.admin.command("ping")
124+
metadata = _get_handshake_driver_info(self.handshake_req)
125+
driver_metadata = metadata["driver"]
126+
name, version, platform = (
127+
driver_metadata["name"],
128+
driver_metadata["version"],
129+
metadata["platform"],
130+
)
131+
# feels like i should do something to check that it is initially sent
132+
self.assertIsNotNone(name)
133+
self.assertIsNotNone(version)
134+
self.assertIsNotNone(platform)
135+
136+
# add data
137+
add_name, add_version, add_platform = "framework", "2.0", "Framework Platform"
138+
client._append_metadata(DriverInfo(add_name, add_version, add_platform))
139+
# check new data isn't sent
140+
self.handshake_req = None
141+
await client.admin.command("ping")
142+
# if it was an actual handshake request, client data would be in the ping request which would start the handshake i think
143+
self.assertNotIn("client", self.handshake_req)
144+
self.assertEqual(listener.event_count(ConnectionClosedEvent), 0)
145+
146+
await client.close()
147+
148+
149+
# THESE ARE MY NOTES TO SELF, PLAESE IGNORE
150+
# two options
151+
# emit events with a flag, so when testing (like now), we can emit more stuff
152+
# or we can mock it? but not sure if that ruins the spirit of the test -> this would be easier tho..
153+
# we'd send, mock server would receive and send back to us?
154+
# use mockup DB!
155+
# usually, create mockupDB instance and then tell it to automatically respond to automatic responses, i'll need to change that for this but i also don't want to mess up SDAM stuff?
156+
# i can get logs from server (in mongo orchestration) if mockup DB is too annoying (maybe get log?)
157+
# ask the team generally but jib thinks we should emit events with a flag
158+
# make sure to state pros and cons for each ticket
159+
160+
if __name__ == "__main__":
161+
unittest.main()

‎test/test_client_metadata.py

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
# Copyright 2013-present MongoDB, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
from __future__ import annotations
15+
16+
import asyncio
17+
import time
18+
import unittest
19+
from test.utils_shared import CMAPListener
20+
from typing import Optional
21+
22+
import pytest
23+
from mockupdb import MockupDB, OpMsgReply
24+
25+
from pymongo import MongoClient
26+
from pymongo.driver_info import DriverInfo
27+
from pymongo.monitoring import ConnectionClosedEvent
28+
29+
pytestmark = pytest.mark.mockupdb
30+
31+
_IS_SYNC = True
32+
33+
34+
def _get_handshake_driver_info(request):
35+
assert "client" in request
36+
return request["client"]
37+
38+
39+
class TestClientMetadataProse(unittest.TestCase):
40+
def setUp(self):
41+
self.server = MockupDB()
42+
self.handshake_req = None
43+
44+
def respond(r):
45+
# Only save the very first request from the driver.
46+
if self.handshake_req is None:
47+
self.handshake_req = r
48+
return r.reply(OpMsgReply(minWireVersion=0, maxWireVersion=13))
49+
50+
self.server.autoresponds(respond)
51+
self.server.run()
52+
self.addCleanup(self.server.stop)
53+
54+
def check_metadata_added(
55+
self, add_name: Optional[str], add_version: Optional[str], add_platform: Optional[str]
56+
) -> None:
57+
client = MongoClient(
58+
"mongodb://" + self.server.address_string,
59+
maxIdleTimeMS=1,
60+
driver=DriverInfo("library", "1.2", "Library Platform"),
61+
)
62+
63+
# send initial metadata
64+
client.admin.command("ping")
65+
metadata = _get_handshake_driver_info(self.handshake_req)
66+
driver_metadata = metadata["driver"]
67+
name, version, platform = (
68+
driver_metadata["name"],
69+
driver_metadata["version"],
70+
metadata["platform"],
71+
)
72+
time.sleep(0.005)
73+
74+
# add data
75+
client._append_metadata(DriverInfo(add_name, add_version, add_platform))
76+
# reset
77+
self.handshake_req = None
78+
client.admin.command("ping")
79+
new_metadata = _get_handshake_driver_info(self.handshake_req)
80+
# compare
81+
self.assertEqual(
82+
new_metadata["driver"]["name"], f"{name}|{add_name}" if add_name is not None else name
83+
)
84+
self.assertEqual(
85+
new_metadata["driver"]["version"],
86+
f"{version}|{add_version}" if add_version is not None else version,
87+
)
88+
self.assertEqual(
89+
new_metadata["platform"],
90+
f"{platform}|{add_platform}" if add_platform is not None else platform,
91+
)
92+
93+
metadata.pop("driver")
94+
metadata.pop("platform")
95+
new_metadata.pop("driver")
96+
new_metadata.pop("platform")
97+
self.assertEqual(metadata, new_metadata)
98+
99+
client.close()
100+
101+
def test_append_metadata(self):
102+
self.check_metadata_added("framework", "2.0", "Framework Platform")
103+
104+
def test_append_metadata_platform_none(self):
105+
self.check_metadata_added("framework", "2.0", None)
106+
107+
def test_append_metadata_version_none(self):
108+
self.check_metadata_added("framework", None, "Framework Platform")
109+
110+
def test_append_platform_version_none(self):
111+
self.check_metadata_added("framework", None, None)
112+
113+
def test_doesnt_update_established_connections(self):
114+
listener = CMAPListener()
115+
client = MongoClient(
116+
"mongodb://" + self.server.address_string,
117+
maxIdleTimeMS=1,
118+
driver=DriverInfo("library", "1.2", "Library Platform"),
119+
event_listeners=[listener],
120+
)
121+
122+
# send initial metadata
123+
client.admin.command("ping")
124+
metadata = _get_handshake_driver_info(self.handshake_req)
125+
driver_metadata = metadata["driver"]
126+
name, version, platform = (
127+
driver_metadata["name"],
128+
driver_metadata["version"],
129+
metadata["platform"],
130+
)
131+
# feels like i should do something to check that it is initially sent
132+
self.assertIsNotNone(name)
133+
self.assertIsNotNone(version)
134+
self.assertIsNotNone(platform)
135+
136+
# add data
137+
add_name, add_version, add_platform = "framework", "2.0", "Framework Platform"
138+
client._append_metadata(DriverInfo(add_name, add_version, add_platform))
139+
# check new data isn't sent
140+
self.handshake_req = None
141+
client.admin.command("ping")
142+
# if it was an actual handshake request, client data would be in the ping request which would start the handshake i think
143+
self.assertNotIn("client", self.handshake_req)
144+
self.assertEqual(listener.event_count(ConnectionClosedEvent), 0)
145+
146+
client.close()
147+
148+
149+
# THESE ARE MY NOTES TO SELF, PLAESE IGNORE
150+
# two options
151+
# emit events with a flag, so when testing (like now), we can emit more stuff
152+
# or we can mock it? but not sure if that ruins the spirit of the test -> this would be easier tho..
153+
# we'd send, mock server would receive and send back to us?
154+
# use mockup DB!
155+
# usually, create mockupDB instance and then tell it to automatically respond to automatic responses, i'll need to change that for this but i also don't want to mess up SDAM stuff?
156+
# i can get logs from server (in mongo orchestration) if mockup DB is too annoying (maybe get log?)
157+
# ask the team generally but jib thinks we should emit events with a flag
158+
# make sure to state pros and cons for each ticket
159+
160+
if __name__ == "__main__":
161+
unittest.main()

‎tools/synchro.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ def async_only_test(f: str) -> bool:
209209
"test_client.py",
210210
"test_client_bulk_write.py",
211211
"test_client_context.py",
212+
"test_client_metadata.py",
212213
"test_collation.py",
213214
"test_collection.py",
214215
"test_collection_management.py",

0 commit comments

Comments
 (0)