diff --git a/backend/src/mirrors_qa_backend/cli/worker.py b/backend/src/mirrors_qa_backend/cli/worker.py index 9cc3359..6377881 100644 --- a/backend/src/mirrors_qa_backend/cli/worker.py +++ b/backend/src/mirrors_qa_backend/cli/worker.py @@ -3,17 +3,17 @@ from mirrors_qa_backend import logger from mirrors_qa_backend.db import Session -from mirrors_qa_backend.db.country import create_country +from mirrors_qa_backend.db.country import update_countries as update_db_countries from mirrors_qa_backend.db.worker import create_worker as create_db_worker from mirrors_qa_backend.db.worker import update_worker as update_db_worker -def get_country_data(*country_codes: str) -> dict[str, str]: +def get_country_mapping(country_codes: list[str]) -> dict[str, str]: """Fetch the country names from the country codes. Maps the country code to the country name. """ - country_data: dict[str, str] = {} + country_mapping: dict[str, str] = {} # Ensure all the countries are valid country codes for country_code in country_codes: if len(country_code) != 2: # noqa: PLR2004 @@ -22,18 +22,20 @@ def get_country_data(*country_codes: str) -> dict[str, str]: ) if country := pycountry.countries.get(alpha_2=country_code): - country_data[country_code] = country.name + country_mapping[country_code] = country.name else: raise ValueError(f"'{country_code}' is not valid country code") - return country_data + return country_mapping -def create_worker(worker_id: str, private_key_data: bytes, country_codes: list[str]): +def create_worker( + worker_id: str, private_key_data: bytes, initial_country_codes: list[str] +): """Create a worker in the DB. Assigns the countries for a worker to run tests from. """ - country_data = get_country_data(*country_codes) + country_mapping = get_country_mapping(initial_country_codes) private_key = serialization.load_pem_private_key( private_key_data, password=None ) # pyright: ignore[reportReturnType] @@ -41,14 +43,11 @@ def create_worker(worker_id: str, private_key_data: bytes, country_codes: list[s with Session.begin() as session: # Update the database with the countries in case those countries don't # exist yet. - for country_code, country_name in country_data.items(): - create_country( - session, country_code=country_code, country_name=country_name - ) + update_db_countries(session, country_mapping) create_db_worker( session, worker_id, - country_codes, + initial_country_codes, private_key, # pyright: ignore [reportGeneralTypeIssues, reportArgumentType] ) @@ -60,10 +59,7 @@ def update_worker(worker_id: str, country_codes: list[str]): Updates the ountries for a worker to run tests from. """ - country_data = get_country_data(*country_codes) + country_mapping = get_country_mapping(country_codes) with Session.begin() as session: - for country_code, country_name in country_data.items(): - create_country( - session, country_code=country_code, country_name=country_name - ) + update_db_countries(session, country_mapping) update_db_worker(session, worker_id, country_codes) diff --git a/backend/src/mirrors_qa_backend/db/country.py b/backend/src/mirrors_qa_backend/db/country.py index 821f3f3..f085c11 100644 --- a/backend/src/mirrors_qa_backend/db/country.py +++ b/backend/src/mirrors_qa_backend/db/country.py @@ -42,3 +42,9 @@ def create_country( .on_conflict_do_nothing(index_elements=["code"]) ) return get_country(session, country_code) + + +def update_countries(session: OrmSession, country_mapping: dict[str, str]) -> None: + """Updates the list of countries in the database.""" + for country_code, country_name in country_mapping.items(): + create_country(session, country_code=country_code, country_name=country_name) diff --git a/dev/docker-compose.yaml b/dev/docker-compose.yaml index d8e97f0..9137b01 100644 --- a/dev/docker-compose.yaml +++ b/dev/docker-compose.yaml @@ -69,7 +69,7 @@ services: - BACKEND_API_URI=http://backend - SLEEP_DURATION=5m - TASK_WORKER_IMAGE=mirrors-qa-task-worker - - TEST_FILE_PATH=/zim/wikipedia/wikipedia_guw_all_nopic_2023-02.zim + - TEST_FILE_PATH=/zim/wikipedia/speedtest_en_blob-mini_2024-05.zim command: mirrors-qa-manager --verbose ${WORKER_ID} task-worker: build: