willtools-demo/WillTools.extension/tests/test_dsh_pool.py

# -*- coding: utf-8 -*-
"""Tests for tools/dsh_pool.py -- the bounded-concurrency fan-out over
tools/dsh_worker.py.

Loaded by path, like test_dsh_worker.py: host CPython, outside
WillTools.extension entirely. `run_task` is faked here the same way
dsh_worker.py's own tests fake `which`/`runner` -- no real dsh install
needed. What matters: input order is preserved in the output regardless of
which task finishes first, one task raising never takes the others down,
a refusal is distinguishable from a short success, and the worker cap is
actually respected rather than merely accepted as an argument.
"""
import json
import os
import sys
import threading
import time

import pytest

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__))))
TOOLS = os.path.join(ROOT, 'tools')
if TOOLS not in sys.path:
    sys.path.insert(0, TOOLS)

dsh_pool = pytest.importorskip('dsh_pool')
dsh_worker = pytest.importorskip('dsh_worker')


# --------------------------------------------------------------- run_pool

def test_run_pool_returns_empty_list_for_no_tasks():
    assert dsh_pool.run_pool([], run_task=lambda *a, **k: 'x') == []


def test_run_pool_preserves_input_order_even_when_they_finish_out_of_order():
    # task i sleeps for (N - i) * a tiny amount, so the LAST task submitted
    # finishes FIRST -- if the pool used completion order instead of the
    # submission index, this would come back scrambled.
    order = ['first', 'second', 'third', 'fourth']

    def fake_run_task(task_text, profile, timeout_s):
        idx = order.index(task_text)
        time.sleep((len(order) - idx) * 0.02)
        return 'reply-for-' + task_text

    results = dsh_pool.run_pool(order, max_workers=4,
                                run_task=fake_run_task)
    assert [r['reply'] for r in results] == \
        ['reply-for-' + t for t in order]


def test_run_pool_marks_one_failure_without_taking_the_batch_down():
    def fake_run_task(task_text, profile, timeout_s):
        if task_text == 'bad':
            raise dsh_worker.DshRefusal('dsh exited 1: boom')
        return 'ok: ' + task_text

    results = dsh_pool.run_pool(['good1', 'bad', 'good2'],
                                run_task=fake_run_task)
    assert results[0] == {'ok': True, 'reply': 'ok: good1'}
    assert results[1]['ok'] is False
    assert results[1]['refused'] is True
    assert 'boom' in results[1]['reason']
    assert results[2] == {'ok': True, 'reply': 'ok: good2'}


def test_run_pool_catches_an_unexpected_exception_too_not_just_dshrefusal():
    def fake_run_task(task_text, profile, timeout_s):
        if task_text == 'crashes':
            raise ValueError('something dsh_worker did not anticipate')
        return 'fine'

    results = dsh_pool.run_pool(['ok1', 'crashes', 'ok2'],
                                run_task=fake_run_task)
    assert results[0]['ok'] is True
    assert results[1]['ok'] is False
    assert results[1]['refused'] is True
    assert 'ValueError' in results[1]['reason']
    assert results[2]['ok'] is True


def test_run_pool_distinguishes_refusal_from_a_short_successful_reply():
    def fake_run_task(task_text, profile, timeout_s):
        if task_text == 'empty-ish':
            raise dsh_worker.DshRefusal(
                'exited 0 but printed no reply -- could not judge, not '
                'judged-empty')
        return 'ok'  # a genuinely short, successful reply

    results = dsh_pool.run_pool(['short-ok', 'empty-ish'],
                                run_task=fake_run_task)
    assert results[0] == {'ok': True, 'reply': 'ok'}
    assert results[1]['ok'] is False
    assert results[1]['refused'] is True
    # a short *successful* reply must never be shaped like a refusal, and
    # vice versa -- the two dicts must not be mistakable for each other.
    assert 'reply' not in results[1]
    assert 'reason' not in results[0]


def test_run_pool_respects_the_max_workers_cap():
    concurrent = [0]
    peak = [0]
    lock = threading.Lock()

    def fake_run_task(task_text, profile, timeout_s):
        with lock:
            concurrent[0] += 1
            peak[0] = max(peak[0], concurrent[0])
        time.sleep(0.05)
        with lock:
            concurrent[0] -= 1
        return 'done'

    dsh_pool.run_pool(['t{0}'.format(i) for i in range(10)],
                      max_workers=2, run_task=fake_run_task)
    assert peak[0] <= 2


def test_run_pool_passes_profile_and_timeout_through():
    seen = []

    def fake_run_task(task_text, profile, timeout_s):
        seen.append((task_text, profile, timeout_s))
        return 'ok'

    dsh_pool.run_pool(['x'], profile='custom-profile', timeout_s=17,
                      run_task=fake_run_task)
    assert seen == [('x', 'custom-profile', 17)]


# --------------------------------------------------------------- _load_tasks

def test_load_tasks_accepts_plain_strings(tmp_path):
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(json.dumps(['a', 'b']), encoding='utf-8')
    texts, extras = dsh_pool._load_tasks(str(tasks_file))
    assert texts == ['a', 'b']
    assert extras == [{}, {}]


def test_load_tasks_accepts_objects_with_a_task_key_and_echoes_extras(
        tmp_path):
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(
        json.dumps([{'id': 'a', 'task': 'do a'},
                   {'id': 'b', 'task': 'do b', 'label': 'B task'}]),
        encoding='utf-8')
    texts, extras = dsh_pool._load_tasks(str(tasks_file))
    assert texts == ['do a', 'do b']
    assert extras == [{'id': 'a'}, {'id': 'b', 'label': 'B task'}]


def test_load_tasks_rejects_a_non_array_file(tmp_path):
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(json.dumps({'not': 'an array'}), encoding='utf-8')
    with pytest.raises(SystemExit):
        dsh_pool._load_tasks(str(tasks_file))


def test_load_tasks_rejects_a_task_item_with_no_task_key(tmp_path):
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(json.dumps([{'id': 'a'}]), encoding='utf-8')
    with pytest.raises(SystemExit):
        dsh_pool._load_tasks(str(tasks_file))


# --------------------------------------------------------------- CLI: main

def test_main_writes_results_to_out_file_preserving_order_and_extras(
        tmp_path, monkeypatch):
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(
        json.dumps([{'id': 'a', 'task': 'do a'},
                   {'id': 'b', 'task': 'do b'}]),
        encoding='utf-8')
    out_file = tmp_path / 'results.json'

    def fake_run_task(task_text, profile, timeout_s):
        return 'reply-for-' + task_text

    monkeypatch.setattr(dsh_worker, 'run_task', fake_run_task)

    code = dsh_pool.main(['--tasks', str(tasks_file), '--out',
                          str(out_file)])
    assert code == 0

    results = json.loads(out_file.read_text(encoding='utf-8'))
    assert results == [
        {'id': 'a', 'ok': True, 'reply': 'reply-for-do a'},
        {'id': 'b', 'ok': True, 'reply': 'reply-for-do b'},
    ]


def test_main_prints_results_to_stdout_when_no_out_given(tmp_path, capsys,
                                                          monkeypatch):
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(json.dumps(['solo task']), encoding='utf-8')

    monkeypatch.setattr(dsh_worker, 'run_task',
                        lambda task_text, profile, timeout_s: 'the reply')

    code = dsh_pool.main(['--tasks', str(tasks_file)])
    assert code == 0
    out = json.loads(capsys.readouterr().out)
    assert out == [{'ok': True, 'reply': 'the reply'}]


def test_main_returns_zero_even_when_some_tasks_are_refused(tmp_path,
                                                             monkeypatch):
    # The batch itself completing is what the exit code reports -- a
    # per-task refusal must never look like the whole run failed.
    tasks_file = tmp_path / 'tasks.json'
    tasks_file.write_text(json.dumps(['ok-task', 'bad-task']),
                          encoding='utf-8')

    def fake_run_task(task_text, profile, timeout_s):
        if task_text == 'bad-task':
            raise dsh_worker.DshRefusal('nope')
        return 'fine'

    monkeypatch.setattr(dsh_worker, 'run_task', fake_run_task)
    code = dsh_pool.main(['--tasks', str(tasks_file)])
    assert code == 0