diff --git a/tests/integration/sharding_tests/helpers.py b/tests/integration/sharding_tests/helpers.py index 711af50a7..299267cd8 100644 --- a/tests/integration/sharding_tests/helpers.py +++ b/tests/integration/sharding_tests/helpers.py @@ -65,3 +65,17 @@ def get_databases_for_shard(mongos_client, shard_name) -> Optional[List[str]]: return return databases_collection.distinct("_id", {"primary": shard_name}) + + +def has_correct_shards(mongos_client, expected_shards: List[str]) -> bool: + """Returns true if the cluster config has the expected shards.""" + shard_names = get_cluster_shards(mongos_client) + return shard_names == set(expected_shards) + + +def shard_has_databases( + mongos_client, shard_name: str, expected_databases_on_shard: List[str] +) -> bool: + """Returns true if the provided shard is a primary for the provided databases.""" + databases_on_shard = get_databases_for_shard(mongos_client, shard_name=shard_name) + return set(databases_on_shard) == set(expected_databases_on_shard) diff --git a/tests/integration/sharding_tests/test_sharding.py b/tests/integration/sharding_tests/test_sharding.py index ccd874264..89cb23c37 100644 --- a/tests/integration/sharding_tests/test_sharding.py +++ b/tests/integration/sharding_tests/test_sharding.py @@ -6,8 +6,8 @@ from .helpers import ( generate_mongodb_client, - get_cluster_shards, - get_databases_for_shard, + has_correct_shards, + shard_has_databases, verify_data_mongodb, write_data_to_mongodb, ) @@ -31,20 +31,23 @@ async def test_build_and_deploy(ops_test: OpsTest) -> None: my_charm = await ops_test.build_charm(".") await ops_test.model.deploy( my_charm, - num_units=2, + num_units=1, config={"role": "config-server"}, application_name=CONFIG_SERVER_APP_NAME, ) await ops_test.model.deploy( - my_charm, num_units=2, config={"role": "shard"}, application_name=SHARD_ONE_APP_NAME + my_charm, num_units=1, config={"role": "shard"}, application_name=SHARD_ONE_APP_NAME ) await ops_test.model.deploy( - my_charm, num_units=2, config={"role": "shard"}, application_name=SHARD_TWO_APP_NAME + my_charm, num_units=1, config={"role": "shard"}, application_name=SHARD_TWO_APP_NAME + ) + await ops_test.model.deploy( + my_charm, num_units=1, config={"role": "shard"}, application_name=SHARD_THREE_APP_NAME ) async with ops_test.fast_forward(): await ops_test.model.wait_for_idle( - apps=[CONFIG_SERVER_APP_NAME, SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME], + apps=[CONFIG_SERVER_APP_NAME, SHARD_ONE_APP_NAME, SHARD_THREE_APP_NAME], idle_period=20, raise_on_blocked=False, timeout=TIMEOUT, @@ -66,24 +69,33 @@ async def test_cluster_active(ops_test: OpsTest) -> None: f"{SHARD_TWO_APP_NAME}:{SHARD_REL_NAME}", f"{CONFIG_SERVER_APP_NAME}:{CONFIG_SERVER_REL_NAME}", ) + await ops_test.model.integrate( + f"{SHARD_THREE_APP_NAME}:{SHARD_REL_NAME}", + f"{CONFIG_SERVER_APP_NAME}:{CONFIG_SERVER_REL_NAME}", + ) async with ops_test.fast_forward(): await ops_test.model.wait_for_idle( - apps=[CONFIG_SERVER_APP_NAME, SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME], + apps=[ + CONFIG_SERVER_APP_NAME, + SHARD_ONE_APP_NAME, + SHARD_TWO_APP_NAME, + SHARD_THREE_APP_NAME, + ], idle_period=20, status="active", timeout=TIMEOUT, raise_on_error=False, ) - # verify sharded cluster config mongos_client = await generate_mongodb_client( ops_test, app_name=CONFIG_SERVER_APP_NAME, mongos=True ) - shard_names = get_cluster_shards(mongos_client) - expected_shard_names = [SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME] - assert shard_names == set( - expected_shard_names + + # verify sharded cluster config + assert has_correct_shards( + mongos_client, + expected_shards=[SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME, SHARD_THREE_APP_NAME], ), "Config server did not process config properly" # TODO Future PR: assert that CONFIG_SERVER_APP_NAME, SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME @@ -98,49 +110,49 @@ async def test_sharding(ops_test: OpsTest) -> None: ops_test, app_name=CONFIG_SERVER_APP_NAME, mongos=True ) - # write data to shard one + # write data to shard two 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) + mongos_client.admin.command("movePrimary", "animals_database_1", to=SHARD_TWO_APP_NAME) - # write data to shard two + # write data to shard three 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) + mongos_client.admin.command("movePrimary", "animals_database_2", to=SHARD_THREE_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 + # log into shard two verify data + shard_two_client = await generate_mongodb_client( + ops_test, app_name=SHARD_TWO_APP_NAME, mongos=False ) has_correct_data = verify_data_mongodb( - shard_one_client, + shard_two_client, db_name="animals_database_1", coll_name="horses", key="horse-breed", value="unicorn", ) - assert has_correct_data, "data not written to shard-one" + assert has_correct_data, "data not written to shard-two" # log into shard 2 verify data - shard_two_client = await generate_mongodb_client( - ops_test, app_name=SHARD_TWO_APP_NAME, mongos=False + shard_three_client = await generate_mongodb_client( + ops_test, app_name=SHARD_THREE_APP_NAME, mongos=False ) has_correct_data = verify_data_mongodb( - shard_two_client, + shard_three_client, db_name="animals_database_2", coll_name="horses", key="horse-breed", value="pegasus", ) - assert has_correct_data, "data not written to shard-two" + assert has_correct_data, "data not written to shard-three" async def test_shard_removal(ops_test: OpsTest) -> None: @@ -151,29 +163,6 @@ async def test_shard_removal(ops_test: OpsTest) -> None: - The balancer is turned back on if turned off. - Config server supp orts removing multiple shards. """ - # add a third shard, so that we can remove two shards at a time. - my_charm = await ops_test.build_charm(".") - await ops_test.model.deploy( - my_charm, num_units=1, config={"role": "shard"}, application_name=SHARD_THREE_APP_NAME - ) - await ops_test.model.integrate( - f"{SHARD_THREE_APP_NAME}:{SHARD_REL_NAME}", - f"{CONFIG_SERVER_APP_NAME}:{CONFIG_SERVER_REL_NAME}", - ) - - await ops_test.model.wait_for_idle( - apps=[ - CONFIG_SERVER_APP_NAME, - SHARD_ONE_APP_NAME, - SHARD_TWO_APP_NAME, - SHARD_THREE_APP_NAME, - ], - idle_period=20, - status="active", - timeout=TIMEOUT, - raise_on_error=False, - ) - # turn off balancer. mongos_client = await generate_mongodb_client( ops_test, app_name=CONFIG_SERVER_APP_NAME, mongos=True @@ -184,17 +173,22 @@ async def test_shard_removal(ops_test: OpsTest) -> None: # remove two shards at the same time await ops_test.model.applications[CONFIG_SERVER_APP_NAME].remove_relation( - f"{SHARD_ONE_APP_NAME}:{SHARD_REL_NAME}", + f"{SHARD_TWO_APP_NAME}:{SHARD_REL_NAME}", f"{CONFIG_SERVER_APP_NAME}:{CONFIG_SERVER_REL_NAME}", ) await ops_test.model.applications[CONFIG_SERVER_APP_NAME].remove_relation( - f"{SHARD_TWO_APP_NAME}:{SHARD_REL_NAME}", + f"{SHARD_THREE_APP_NAME}:{SHARD_REL_NAME}", f"{CONFIG_SERVER_APP_NAME}:{CONFIG_SERVER_REL_NAME}", ) async with ops_test.fast_forward(): await ops_test.model.wait_for_idle( - apps=[CONFIG_SERVER_APP_NAME, SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME], + apps=[ + CONFIG_SERVER_APP_NAME, + SHARD_ONE_APP_NAME, + SHARD_TWO_APP_NAME, + SHARD_THREE_APP_NAME, + ], idle_period=20, status="active", timeout=TIMEOUT, @@ -207,19 +201,16 @@ async def test_shard_removal(ops_test: OpsTest) -> None: balancer_state = mongos_client.admin.command("balancerStatus") assert balancer_state["mode"] != "off", "balancer not turned back on from config server" - # veriy sharded cluster config - shard_names = get_cluster_shards(mongos_client) - expected_shard_names = [SHARD_THREE_APP_NAME] - assert shard_names == set( - expected_shard_names + # verify sharded cluster config + assert has_correct_shards( + mongos_client, expected_shards=[SHARD_ONE_APP_NAME] ), "Config server did not process config properly" - # verify databases that had primaries shard-one and shard-two are now on shard-three - databases_on_shard = get_databases_for_shard(mongos_client, shard_name=SHARD_THREE_APP_NAME) - expected_databases_on_shard = ["animals_database_1", "animals_database_2"] - assert databases_on_shard, "No databases on the final shard." - assert set(databases_on_shard) == set( - expected_databases_on_shard + # verify no data lost + assert shard_has_databases( + mongos_client, + shard_name=SHARD_ONE_APP_NAME, + expected_databases_on_shard=["animals_database_1", "animals_database_2"], ), "Not all databases on final shard" @@ -258,6 +249,22 @@ async def test_removal_of_non_primary_shard(ops_test: OpsTest): raise_on_error=False, ) + mongos_client = await generate_mongodb_client( + ops_test, app_name=CONFIG_SERVER_APP_NAME, mongos=True + ) + + # verify sharded cluster config + assert has_correct_shards( + mongos_client, expected_shards=[SHARD_ONE_APP_NAME] + ), "Config server did not process config properly" + + # verify no data lost + assert shard_has_databases( + mongos_client, + shard_name=SHARD_ONE_APP_NAME, + expected_databases_on_shard=["animals_database_1", "animals_database_2"], + ), "Not all databases on final shard" + async def test_unconventual_shard_removal(ops_test: OpsTest): """Tests that removing a shard application safely drains data. @@ -273,7 +280,7 @@ async def test_unconventual_shard_removal(ops_test: OpsTest): async with ops_test.fast_forward(): await ops_test.model.wait_for_idle( - apps=[CONFIG_SERVER_APP_NAME, SHARD_TWO_APP_NAME, SHARD_THREE_APP_NAME], + apps=[CONFIG_SERVER_APP_NAME, SHARD_ONE_APP_NAME, SHARD_TWO_APP_NAME], idle_period=20, status="active", timeout=TIMEOUT, @@ -282,20 +289,18 @@ async def test_unconventual_shard_removal(ops_test: OpsTest): ops_test.model.remove_application(SHARD_TWO_APP_NAME, block_until_done=True) - # veriy sharded cluster config mongos_client = await generate_mongodb_client( ops_test, app_name=CONFIG_SERVER_APP_NAME, mongos=True ) - shard_names = get_cluster_shards(mongos_client) - expected_shard_names = [SHARD_THREE_APP_NAME] - assert shard_names == set( - expected_shard_names + + # verify sharded cluster config + assert has_correct_shards( + mongos_client, expected_shards=[SHARD_ONE_APP_NAME] ), "Config server did not process config properly" # verify no data lost - databases_on_shard = get_databases_for_shard(mongos_client, shard_name=SHARD_THREE_APP_NAME) - expected_databases_on_shard = ["animals_database_1", "animals_database_2"] - assert databases_on_shard, "No databases on the final shard." - assert set(databases_on_shard) == set( - expected_databases_on_shard + assert shard_has_databases( + mongos_client, + shard_name=SHARD_ONE_APP_NAME, + expected_databases_on_shard=["animals_database_1", "animals_database_2"], ), "Not all databases on final shard"