# -*- coding:utf-8 -*- import random import textwrap from pathlib import Path import pytest from dw_base.datax.worker import load_workers_ini, select_worker @pytest.fixture def pool_ini(tmp_path: Path) -> str: p = tmp_path / 'workers.ini' p.write_text(textwrap.dedent('''\ [release] host = cdhmaster02 [weights] cdhmaster02 = 1 cdhnode01 = 3 cdhnode02 = 3 cdhnode03 = 3 ''')) return str(p) def test_load_workers_ini(pool_ini): pool = load_workers_ini(pool_ini) assert pool.release_host == 'cdhmaster02' assert pool.weights == {'cdhmaster02': 1, 'cdhnode01': 3, 'cdhnode02': 3, 'cdhnode03': 3} assert len(pool.queue) == 10 assert pool.queue.count('cdhnode01') == 3 assert pool.queue.count('cdhmaster02') == 1 def test_select_worker_non_release_user_falls_back_local(pool_ini): pool = load_workers_ini(pool_ini) assert select_worker(pool, is_release_user=False, is_in_release_dir=True, current_host='dev-box', use_random=True) == 'dev-box' def test_select_worker_outside_release_dir_falls_back_local(pool_ini): pool = load_workers_ini(pool_ini) assert select_worker(pool, is_release_user=True, is_in_release_dir=False, current_host='cdhmaster02', use_random=True) == 'cdhmaster02' def test_select_worker_explicit_host_overrides_random(pool_ini): pool = load_workers_ini(pool_ini) assert select_worker(pool, is_release_user=True, is_in_release_dir=True, current_host='cdhmaster02', host='cdhnode02', use_random=True) == 'cdhnode02' def test_select_worker_random_pick_in_queue(pool_ini): pool = load_workers_ini(pool_ini) rand = random.Random(42) pick = select_worker(pool, is_release_user=True, is_in_release_dir=True, current_host='cdhmaster02', use_random=True, rand=rand) assert pick in pool.queue def test_select_worker_nothing_specified_falls_back_local(pool_ini): pool = load_workers_ini(pool_ini) assert select_worker(pool, is_release_user=True, is_in_release_dir=True, current_host='cdhmaster02') == 'cdhmaster02'