poyee-data-warehouse/tests/unit/datax/test_worker.py

# -*- 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'