|
|
@@ -0,0 +1,66 @@
|
|
|
+# -*- 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'
|