diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 2670a0757..3c0bdab19 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -81,6 +81,7 @@ jobs: - tls-integration - backup-integration - metric-integration + - sharding-integration name: ${{ matrix.tox-environments }} needs: - lint diff --git a/tests/integration/sharding_tests/helpers.py b/tests/integration/sharding_tests/helpers.py index 7dc150194..6b3c14887 100644 --- a/tests/integration/sharding_tests/helpers.py +++ b/tests/integration/sharding_tests/helpers.py @@ -4,6 +4,7 @@ from urllib.parse import quote_plus from pytest_operator.plugin import OpsTest +from pymongo import MongoClient from ..helpers import get_password @@ -11,8 +12,8 @@ MONGOD_PORT = 27017 -async def generate_connection_string(ops_test: OpsTest, app_name: str, mongos: bool) -> str: - """Returns a connection string for mongos/mongod to the provided application.""" +async def generate_mongodb_client(ops_test: OpsTest, app_name: str, mongos: bool): + """Returns a MongoDB client for mongos/mongod.""" hosts = [unit.public_address for unit in ops_test.model.applications[app_name].units] password = await get_password(ops_test, app_name) port = MONGOS_PORT if mongos else MONGOD_PORT @@ -21,9 +22,24 @@ async def generate_connection_string(ops_test: OpsTest, app_name: str, mongos: b auth_source = "" database = "admin" - return ( + return MongoClient( f"mongodb://operator:" f"{quote_plus(password)}@" f"{hosts}/{quote_plus(database)}?" f"{auth_source}" ) + + +def write_data_to_mongodb(client, db_name, coll_name, content) -> None: + """Writes data to the provided collection and database.""" + db = client[db_name] + horses_collection = db[coll_name] + horses_collection.insert_one(content) + + +def verify_data_mongodb(client, db_name, coll_name, key, value) -> bool: + """Checks a key/value pair for a provided collection and database.""" + db = client[db_name] + test_collection = db[coll_name] + query = test_collection.find({}, {key: 1}) + return query[0][key] == value diff --git a/tests/integration/sharding_tests/test_sharding.py b/tests/integration/sharding_tests/test_sharding.py index 9f268e4d7..728b4ffb9 100644 --- a/tests/integration/sharding_tests/test_sharding.py +++ b/tests/integration/sharding_tests/test_sharding.py @@ -2,10 +2,9 @@ # Copyright 2023 Canonical Ltd. # See LICENSE file for licensing details. import pytest -from pymongo import MongoClient from pytest_operator.plugin import OpsTest -from .helpers import generate_connection_string +from .helpers import generate_mongodb_client, write_data_to_mongodb, verify_data_mongodb SHARD_ONE_APP_NAME = "shard-one" SHARD_TWO_APP_NAME = "shard-two" @@ -67,40 +66,51 @@ async def test_cluster_active(ops_test: OpsTest) -> None: async def test_sharding(ops_test: OpsTest) -> None: """Tests writing data to mongos gets propagated to shards.""" - # write data to mongos router - mongos_connection_string = await generate_connection_string( + # write data to mongos on both shards. + mongos_client = await generate_mongodb_client( ops_test, app_name=CONFIG_SERVER_APP_NAME, mongos=True ) - client = MongoClient(mongos_connection_string) - - db = client["animals_database_1"] - horses_collection = db["horses"] - unicorn = {"horse-breed": "unicorn", "real": True} - horses_collection.insert_one(unicorn) - client.admin.command("movePrimary", "animals_database_1", to=SHARD_ONE_APP_NAME) - - db = client["animals_database_2"] - horses_collection = db["horses"] - unicorn = {"horse-breed": "pegasus", "real": True} - horses_collection.insert_one(unicorn) - client.admin.command("movePrimary", "animals_database_2", to=SHARD_TWO_APP_NAME) - - # log into shard 1 verify its presence - shard_one_connection_string = await generate_connection_string( + + # write data to shard one + write_data_to_mongodb( + mongos_client, + db_name="animals_database_1", + coll_name="horses", + content={"horse-breed": "unicorn", "real": True}, + ) + mongos_client.admin.command("movePrimary", "animals_database_1", to=SHARD_ONE_APP_NAME) + + # write data to shard two + write_data_to_mongodb( + mongos_client, + db_name="animals_database_2", + coll_name="horses", + content={"horse-breed": "pegasus", "real": True}, + ) + mongos_client.admin.command("movePrimary", "animals_database_2", to=SHARD_TWO_APP_NAME) + + # log into shard 1 verify data + shard_one_client = await generate_mongodb_client( ops_test, app_name=SHARD_ONE_APP_NAME, mongos=False ) - client = MongoClient(shard_one_connection_string) - db = client["animals_database_1"] - test_collection = db["horses"] - query = test_collection.find({}, {"horse-breed": 1}) - assert query[0]["horse-breed"] == "unicorn", "data not written to shard-one" - - # log into shard 2 verify its presence - shard_two_connection_string = await generate_connection_string( + has_correct_data = verify_data_mongodb( + shard_one_client, + db_name="animals_database_1", + coll_name="horses", + key="horse-breed", + value="unicorn", + ) + assert has_correct_data, "data not written to shard-one" + + # log into shard 2 verify data + shard_two_client = await generate_mongodb_client( ops_test, app_name=SHARD_TWO_APP_NAME, mongos=False ) - client = MongoClient(shard_two_connection_string) - db = client["animals_database_2"] - test_collection = db["horses"] - query = test_collection.find({}, {"horse-breed": 1}) - assert query[0]["horse-breed"] == "pegasus", "data not written to shard-two" + has_correct_data = verify_data_mongodb( + shard_two_client, + db_name="animals_database_2", + coll_name="horses", + key="horse-breed", + value="pegasus", + ) + assert has_correct_data, "data not written to shard-two"