restructure tests

.github/workflows/ci.yaml

@@ -81,6 +81,7 @@ jobs:
           - tls-integration
           - backup-integration
           - metric-integration
+          - sharding-integration
     name: ${{ matrix.tox-environments }}
     needs:
       - lint

tests/integration/sharding_tests/helpers.py

@@ -4,15 +4,16 @@
 from urllib.parse import quote_plus
 
 from pytest_operator.plugin import OpsTest
+from pymongo import MongoClient
 
 from ..helpers import get_password
 
 MONGOS_PORT = 27018
 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

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"