883 lines
34 KiB
Python
883 lines
34 KiB
Python
# -*- encoding=utf8 -*-
|
||
"""
|
||
Master-Worker 本地模拟测试
|
||
|
||
测试策略:
|
||
1. 直接调用模式 (TestDirectDispatcher): 直接调使用 RedisTaskDispatcher 方法模拟Worker行为
|
||
2. PubSub集成测试 (TestPubSubIntegration): 启动 DispatcherService + 模拟Worker线程
|
||
3. 使用本地 Redis (localhost:6379) DB 15 隔离测试数据
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import json
|
||
import time
|
||
import threading
|
||
import tempfile
|
||
import shutil
|
||
import csv
|
||
|
||
import pytest
|
||
import redis
|
||
|
||
# 将项目根目录加入 sys.path，以便导入项目模块
|
||
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
sys.path.insert(0, PROJECT_ROOT)
|
||
|
||
# ==================== 测试用配置覆盖 ====================
|
||
# 在导入项目模块前，先覆盖 config 中的值以使用本地Redis和测试文件
|
||
import config
|
||
config.REDIS_HOST = 'localhost'
|
||
config.REDIS_PORT = 6379
|
||
config.REDIS_MAX_CONNECTIONS_DISPATCHER = 20
|
||
config.TASK_TIMEOUT = 5 # 测试用: 5秒超时（加速测试）
|
||
config.MAX_RETRY_COUNT = 2 # 测试用: 最多重试2次
|
||
config.WORKER_STALE_TIMEOUT = 3 # 测试用: 3秒Worker超时
|
||
|
||
from redis_task_distribute import RedisTaskDispatcher
|
||
from dispatcher_main import DispatcherService
|
||
|
||
# 测试用CSV路径
|
||
TEST_CSV_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'test_task_list.csv')
|
||
|
||
# 使用 Redis DB 15 隔离测试数据
|
||
TEST_REDIS_DB = 15
|
||
|
||
|
||
class RedisTestHelper:
|
||
"""Redis测试辅助类：管理DB 15的连接和清理"""
|
||
|
||
def __init__(self):
|
||
self.redis = redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True)
|
||
|
||
def flush(self):
|
||
"""清空DB 15中所有数据"""
|
||
self.redis.flushdb()
|
||
|
||
def is_available(self):
|
||
"""检查Redis是否可用"""
|
||
try:
|
||
return self.redis.ping()
|
||
except redis.ConnectionError:
|
||
return False
|
||
|
||
|
||
def create_test_dispatcher():
|
||
"""创建使用DB 15的测试Dispatcher
|
||
|
||
通过修改连接池使其连接到 DB 15，与生产数据隔离
|
||
"""
|
||
# 强制重置连接池，确保使用新的配置
|
||
RedisTaskDispatcher._connection_pool = None
|
||
|
||
dispatcher = RedisTaskDispatcher(redis_host='localhost', redis_port=6379)
|
||
# 替换redis连接为DB 15
|
||
dispatcher.redis = redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True)
|
||
return dispatcher
|
||
|
||
|
||
def create_test_service():
|
||
"""创建使用DB 15的测试DispatcherService"""
|
||
# 重置连接池
|
||
RedisTaskDispatcher._connection_pool = None
|
||
|
||
service = DispatcherService(redis_host='localhost', redis_port=6379)
|
||
# 替换redis连接为DB 15
|
||
service.dispatcher.redis = redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True)
|
||
service.redis = service.dispatcher.redis
|
||
return service
|
||
|
||
|
||
# ==================== 跳过条件 ====================
|
||
helper = RedisTestHelper()
|
||
redis_available = helper.is_available()
|
||
skip_no_redis = pytest.mark.skipif(not redis_available, reason="本地Redis服务未运行")
|
||
|
||
|
||
# ==================== 直接调用模式测试 ====================
|
||
@skip_no_redis
|
||
class TestDirectDispatcher:
|
||
"""直接调用 RedisTaskDispatcher 方法模拟各种Worker场景"""
|
||
|
||
def setup_method(self):
|
||
"""每个测试前清空DB 15并创建新的dispatcher"""
|
||
helper.flush()
|
||
self.dispatcher = create_test_dispatcher()
|
||
# 使用临时目录存放CSV输出
|
||
self.temp_dir = tempfile.mkdtemp()
|
||
config.FAILED_TASKS_CSV = os.path.join(self.temp_dir, 'failed_tasks.csv')
|
||
config.SUCCESS_TASKS_CSV = os.path.join(self.temp_dir, 'success_tasks.csv')
|
||
|
||
def teardown_method(self):
|
||
"""每个测试后清理"""
|
||
helper.flush()
|
||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||
|
||
def _load_test_tasks(self):
|
||
"""加载测试任务"""
|
||
count = self.dispatcher.load_tasks_from_csv(TEST_CSV_PATH)
|
||
assert count == 10, f"预期加载10个任务，实际加载了{count}个"
|
||
return count
|
||
|
||
def test_country_codes_are_parsed_and_delivered(self):
|
||
fd, csv_path = tempfile.mkstemp(dir=PROJECT_ROOT, suffix='.csv')
|
||
os.close(fd)
|
||
try:
|
||
with open(csv_path, 'w', newline='', encoding='utf-8') as handle:
|
||
writer = csv.DictWriter(handle, fieldnames=['app_name', 'package_name', 'country_code', 'device_type'])
|
||
writer.writeheader()
|
||
writer.writerow({
|
||
'app_name': 'Country App',
|
||
'package_name': 'com.test.country',
|
||
'country_code': 'cn, us,JP,cn',
|
||
'device_type': '0',
|
||
})
|
||
|
||
count = self.dispatcher.load_tasks_from_csv(csv_path)
|
||
assert count == 1
|
||
|
||
task_key = 'Country App_com.test.country'
|
||
task_details = json.loads(self.dispatcher.redis.hget('task:details', task_key))
|
||
assert task_details['country_code'] == 'cn, us,JP,cn'
|
||
assert task_details['country_codes'] == ['CN', 'US', 'JP']
|
||
|
||
task = self.dispatcher.worker_init(
|
||
worker_id='192.168.1.1_AA:BB:CC:DD:EE:01',
|
||
ip_address='192.168.1.1',
|
||
mac_address='AA:BB:CC:DD:EE:01',
|
||
hostname='test-country-worker',
|
||
platform='Windows',
|
||
)
|
||
assert task is not None
|
||
assert task['country_code'] == 'cn, us,JP,cn'
|
||
assert task['country_codes'] == ['CN', 'US', 'JP']
|
||
finally:
|
||
if os.path.exists(csv_path):
|
||
os.remove(csv_path)
|
||
|
||
# -------- 场景1: 正常流程 --------
|
||
def test_normal_flow(self):
|
||
"""多Worker并行领取任务 → 上报success → 领取下一个 → 直到队列耗尽"""
|
||
self._load_test_tasks()
|
||
|
||
# 创建3个Worker
|
||
workers = {}
|
||
for i in range(3):
|
||
worker_id = f"192.168.1.{i+1}_AA:BB:CC:DD:EE:0{i}"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address=f"192.168.1.{i+1}",
|
||
mac_address=f"AA:BB:CC:DD:EE:0{i}",
|
||
hostname=f"test-pc-{i}",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None, f"Worker {i} 应该能领取到任务"
|
||
workers[worker_id] = task
|
||
|
||
# 验证3个Worker各领到不同任务
|
||
task_keys = [t['task_key'] for t in workers.values()]
|
||
assert len(set(task_keys)) == 3, "3个Worker应领取到3个不同的任务"
|
||
|
||
# 统计信息
|
||
stats = self.dispatcher.get_statistics()
|
||
assert stats['pending'] == 7, "应剩余7个待分发任务"
|
||
assert stats['running'] == 3, "应有3个运行中任务"
|
||
|
||
# 每个Worker循环: 上报成功 → 领取下一个
|
||
completed_count = 3 # 已领取了3个
|
||
for worker_id in list(workers.keys()):
|
||
while True:
|
||
current_task = workers[worker_id]
|
||
next_task = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=current_task['task_key'],
|
||
status='success'
|
||
)
|
||
completed_count += 0 # 上报时才算完成
|
||
if next_task is None:
|
||
break
|
||
workers[worker_id] = next_task
|
||
|
||
# 最终统计
|
||
stats = self.dispatcher.get_statistics()
|
||
assert stats['pending'] == 0, "所有任务应分发完毕"
|
||
assert stats['completed'] == 10, f"所有10个任务应标记为完成，实际: {stats['completed']}"
|
||
assert stats['failed'] == 0, "不应有失败任务"
|
||
|
||
# -------- 场景2: 任务失败重试 --------
|
||
def test_failed_retry(self):
|
||
"""Worker上报failed → 任务重新入队 → 可被其他Worker领取"""
|
||
self._load_test_tasks()
|
||
|
||
# Worker A 领取任务
|
||
worker_a = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task_a = self.dispatcher.worker_init(
|
||
worker_id=worker_a,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-a",
|
||
platform="Windows"
|
||
)
|
||
assert task_a is not None
|
||
failed_task_key = task_a['task_key']
|
||
|
||
# Worker A 上报失败
|
||
next_task = self.dispatcher.worker_report(
|
||
worker_id=worker_a,
|
||
previous_task_key=failed_task_key,
|
||
status='failed',
|
||
message='下载APK失败'
|
||
)
|
||
# 失败后应能领取新任务
|
||
assert next_task is not None, "上报failed后应自动领取下一个任务"
|
||
|
||
# 验证失败任务已重新入队（队列中应该包含它）
|
||
queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1)
|
||
assert failed_task_key in queue_items, f"失败的任务 {failed_task_key} 应重新入队"
|
||
|
||
# Worker B 领取任务，最终能领到之前失败的任务
|
||
worker_b = "192.168.1.2_AA:BB:CC:DD:EE:02"
|
||
task_b = self.dispatcher.worker_init(
|
||
worker_id=worker_b,
|
||
ip_address="192.168.1.2",
|
||
mac_address="AA:BB:CC:DD:EE:02",
|
||
hostname="test-pc-b",
|
||
platform="Windows"
|
||
)
|
||
# Worker B 连续领取，直到找到失败过的任务
|
||
found = False
|
||
checked_tasks = [task_b['task_key']] if task_b else []
|
||
while task_b:
|
||
if task_b['task_key'] == failed_task_key:
|
||
found = True
|
||
break
|
||
task_b = self.dispatcher.worker_report(
|
||
worker_id=worker_b,
|
||
previous_task_key=task_b['task_key'],
|
||
status='success'
|
||
)
|
||
if task_b:
|
||
checked_tasks.append(task_b['task_key'])
|
||
|
||
assert found, f"Worker B 应能领到之前失败的任务 {failed_task_key}，已检查: {checked_tasks}"
|
||
|
||
# -------- 场景3: Worker stop --------
|
||
def test_worker_stop(self):
|
||
"""Worker上报stop → 从列表移除 → 不再分配任务"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-stop",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None
|
||
|
||
# 上报stop，附带告警消息
|
||
result = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=task['task_key'],
|
||
status='stop',
|
||
message='设备ADB连接异常，无法恢复'
|
||
)
|
||
|
||
# stop后不应返回新任务
|
||
assert result is None, "stop状态不应返回新任务"
|
||
|
||
# Worker应从列表中移除
|
||
workers = self.dispatcher.get_registered_workers()
|
||
worker_ids = [w['worker_id'] for w in workers]
|
||
assert worker_id not in worker_ids, "stop后Worker应从列表移除"
|
||
|
||
# 该Worker再次上报应返回错误码
|
||
error_result = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key="any_task",
|
||
status='success'
|
||
)
|
||
assert error_result is not None, "已移除的Worker上报应返回错误码"
|
||
assert error_result.get('error_code') == -1, "错误码应为-1"
|
||
|
||
# -------- 场景4: 重复初始化 --------
|
||
def test_duplicate_init(self):
|
||
"""Worker重新调用init → 应返回之前分配的未完成任务"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
|
||
# 第一次初始化
|
||
task1 = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-dup",
|
||
platform="Windows"
|
||
)
|
||
assert task1 is not None
|
||
original_task_key = task1['task_key']
|
||
|
||
# 第二次初始化（模拟Worker重启）
|
||
task2 = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-dup",
|
||
platform="Windows"
|
||
)
|
||
assert task2 is not None, "重复初始化应返回任务"
|
||
assert task2['task_key'] == original_task_key, \
|
||
f"重复初始化应返回之前的任务 {original_task_key}，实际: {task2['task_key']}"
|
||
|
||
# -------- 场景5: Worker超时清理 --------
|
||
def test_stale_worker_cleanup(self):
|
||
"""模拟Worker长时间无心跳 → cleanup_stale_workers清理 → 任务重新入队"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-stale",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None
|
||
stale_task_key = task['task_key']
|
||
|
||
# 模拟Worker超时：手动修改 last_update 为很久以前
|
||
worker_info_json = self.dispatcher.redis.hget("workers:info", worker_id)
|
||
worker_info = json.loads(worker_info_json)
|
||
worker_info['last_update'] = time.time() - 10000 # 10000秒前
|
||
self.dispatcher.redis.hset("workers:info", worker_id, json.dumps(worker_info))
|
||
|
||
# 执行清理
|
||
cleaned = self.dispatcher.cleanup_stale_workers(timeout=5) # 5秒超时
|
||
assert cleaned == 1, f"应清理1个Worker，实际清理了{cleaned}个"
|
||
|
||
# 验证Worker已被移除
|
||
workers = self.dispatcher.get_registered_workers()
|
||
assert len(workers) == 0, "超时Worker应被移除"
|
||
|
||
# 验证超时Worker的任务记录了失败（任务被 _record_failed_task 处理）
|
||
status_json = self.dispatcher.redis.hget("task:status", stale_task_key)
|
||
assert status_json is not None, "超时任务应有状态记录"
|
||
status = json.loads(status_json)
|
||
assert status['status'] == 'failed', f"超时任务状态应为failed，实际: {status['status']}"
|
||
|
||
# 验证任务已重新入队（因为retry_count < MAX_RETRY_COUNT）
|
||
queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1)
|
||
assert stale_task_key in queue_items, "超时的任务应重新入队"
|
||
|
||
# -------- 场景6: 任务超时 --------
|
||
def test_task_timeout(self):
|
||
"""模拟任务长时间running → check_timeout → 自动重新入队"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-timeout",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None
|
||
timeout_task_key = task['task_key']
|
||
|
||
# 模拟任务超时：修改start_time和last_retry为很久以前
|
||
status_json = self.dispatcher.redis.hget("task:status", timeout_task_key)
|
||
status = json.loads(status_json)
|
||
status['start_time'] = time.time() - 10000
|
||
status['last_retry'] = time.time() - 10000
|
||
self.dispatcher.redis.hset("task:status", timeout_task_key, json.dumps(status))
|
||
|
||
# 执行超时检查
|
||
timeout_count = self.dispatcher.check_timeout()
|
||
assert timeout_count == 1, f"应检测到1个超时任务，实际: {timeout_count}"
|
||
|
||
# 验证任务已重新入队
|
||
queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1)
|
||
assert timeout_task_key in queue_items, "超时任务应重新入队"
|
||
|
||
# 验证Worker任务映射已清除（Worker可以领新任务）
|
||
current_task = self.dispatcher.redis.hget("worker:tasks", worker_id)
|
||
assert current_task is None, "超时后Worker的任务映射应被清除"
|
||
|
||
# -------- 场景7: 重试达上限 --------
|
||
def test_max_retry_exceeded(self):
|
||
"""任务连续失败超过MAX_RETRY_COUNT → 标记为永久失败"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-maxretry",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None
|
||
doomed_task_key = task['task_key']
|
||
|
||
# MAX_RETRY_COUNT = 2，需要连续失败2次达到上限
|
||
for i in range(config.MAX_RETRY_COUNT):
|
||
# 上报失败
|
||
next_task = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=doomed_task_key,
|
||
status='failed',
|
||
message=f'第{i+1}次失败'
|
||
)
|
||
|
||
if i < config.MAX_RETRY_COUNT - 1:
|
||
# 还没到上限，任务应该重新入队
|
||
# 消费掉领到的新任务，然后继续让worker去领那个失败任务
|
||
# Worker自动领了新任务，需要先完成它再去拿失败重入的任务
|
||
while next_task and next_task['task_key'] != doomed_task_key:
|
||
next_task2 = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=next_task['task_key'],
|
||
status='success'
|
||
)
|
||
next_task = next_task2
|
||
|
||
if next_task is None:
|
||
# 队列中可能还有那个失败任务，手动领取
|
||
# 需要清除worker当前任务映射
|
||
self.dispatcher.redis.hdel("worker:tasks", worker_id)
|
||
self.dispatcher._update_worker_status(worker_id, 'idle', None)
|
||
# 手动分配
|
||
next_task = self.dispatcher._assign_task(worker_id)
|
||
|
||
if next_task and next_task['task_key'] == doomed_task_key:
|
||
continue
|
||
|
||
# 验证任务已标记为永久失败
|
||
failed_tasks = self.dispatcher.redis.smembers("task:failed")
|
||
assert doomed_task_key in failed_tasks, \
|
||
f"任务 {doomed_task_key} 应在永久失败集合中，当前集合: {failed_tasks}"
|
||
|
||
# 验证任务不再入队
|
||
queue_items = self.dispatcher.redis.lrange("task:queue", 0, -1)
|
||
assert doomed_task_key not in queue_items, "永久失败的任务不应在队列中"
|
||
|
||
# -------- 场景8: 并发安全 --------
|
||
def test_concurrent_workers(self):
|
||
"""多线程同时调用worker_init和worker_report，验证不会重复分配"""
|
||
self._load_test_tasks()
|
||
|
||
results = {}
|
||
errors = []
|
||
|
||
def worker_thread(worker_idx):
|
||
"""模拟单个Worker的工作线程"""
|
||
try:
|
||
worker_id = f"192.168.1.{worker_idx}_AA:BB:CC:DD:EE:{worker_idx:02d}"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address=f"192.168.1.{worker_idx}",
|
||
mac_address=f"AA:BB:CC:DD:EE:{worker_idx:02d}",
|
||
hostname=f"test-pc-{worker_idx}",
|
||
platform="Windows"
|
||
)
|
||
|
||
tasks_done = []
|
||
while task:
|
||
tasks_done.append(task['task_key'])
|
||
# 模拟短暂处理
|
||
time.sleep(0.1)
|
||
task = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=tasks_done[-1],
|
||
status='success'
|
||
)
|
||
|
||
results[worker_id] = tasks_done
|
||
except Exception as e:
|
||
errors.append(f"Worker {worker_idx}: {e}")
|
||
|
||
# 启动5个并发Worker
|
||
threads = []
|
||
for i in range(1, 6):
|
||
t = threading.Thread(target=worker_thread, args=(i,))
|
||
threads.append(t)
|
||
t.start()
|
||
|
||
# 等待所有线程完成
|
||
for t in threads:
|
||
t.join(timeout=30)
|
||
|
||
# 验证
|
||
assert len(errors) == 0, f"不应有错误: {errors}"
|
||
|
||
# 收集所有完成的任务
|
||
all_tasks = []
|
||
for tasks in results.values():
|
||
all_tasks.extend(tasks)
|
||
|
||
# 排除重试场景，每个任务应该只被一个Worker完成一次
|
||
assert len(all_tasks) == len(set(all_tasks)), \
|
||
f"不应有重复分配的任务，总数: {len(all_tasks)}，去重后: {len(set(all_tasks))}"
|
||
assert len(all_tasks) == 10, f"所有10个任务应被完成，实际完成: {len(all_tasks)}"
|
||
|
||
# -------- 场景9: 空队列 --------
|
||
def test_empty_queue(self):
|
||
"""队列为空时Worker初始化应返回None"""
|
||
# 不加载任何任务
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-empty",
|
||
platform="Windows"
|
||
)
|
||
assert task is None, "空队列时应返回None"
|
||
|
||
# Worker应已注册但状态为idle
|
||
workers = self.dispatcher.get_registered_workers()
|
||
assert len(workers) == 1, "Worker应已注册"
|
||
assert workers[0]['status'] == 'idle', "Worker应为idle状态"
|
||
|
||
# -------- 场景10: worker_retry重新计时 --------
|
||
def test_worker_retry(self):
|
||
"""Worker调用retry → 任务重新计时"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-retry",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None
|
||
task_key = task['task_key']
|
||
|
||
# 记录原始时间
|
||
status_before = json.loads(self.dispatcher.redis.hget("task:status", task_key))
|
||
original_retry_time = status_before.get('last_retry')
|
||
|
||
# 等一小段时间
|
||
time.sleep(0.5)
|
||
|
||
# 调用retry
|
||
success = self.dispatcher.worker_retry(worker_id, task_key)
|
||
assert success is True, "retry应成功"
|
||
|
||
# 验证时间已更新
|
||
status_after = json.loads(self.dispatcher.redis.hget("task:status", task_key))
|
||
new_retry_time = status_after.get('last_retry')
|
||
assert new_retry_time > original_retry_time, "retry后时间应更新"
|
||
|
||
# -------- 场景11: 无效status验证 --------
|
||
def test_invalid_status(self):
|
||
"""上报无效的status值 → 应返回None"""
|
||
self._load_test_tasks()
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-invalid",
|
||
platform="Windows"
|
||
)
|
||
assert task is not None
|
||
|
||
# 上报无效状态
|
||
result = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=task['task_key'],
|
||
status='invalid_status'
|
||
)
|
||
assert result is None, "无效status应返回None"
|
||
|
||
# 上报非字符串状态
|
||
result2 = self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=task['task_key'],
|
||
status=123
|
||
)
|
||
assert result2 is None, "非字符串status应返回None"
|
||
|
||
# -------- 场景12: 统计信息一致性 --------
|
||
def test_statistics_consistency(self):
|
||
"""验证统计信息在各操作后的一致性"""
|
||
self._load_test_tasks()
|
||
|
||
# 初始状态
|
||
stats = self.dispatcher.get_statistics()
|
||
assert stats['pending'] == 10
|
||
assert stats['running'] == 0
|
||
assert stats['completed'] == 0
|
||
assert stats['failed'] == 0
|
||
|
||
# Worker领取任务后
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
task = self.dispatcher.worker_init(
|
||
worker_id=worker_id,
|
||
ip_address="192.168.1.1",
|
||
mac_address="AA:BB:CC:DD:EE:01",
|
||
hostname="test-pc-stats",
|
||
platform="Windows"
|
||
)
|
||
stats = self.dispatcher.get_statistics()
|
||
assert stats['pending'] == 9
|
||
assert stats['running'] == 1
|
||
|
||
# 上报成功后
|
||
self.dispatcher.worker_report(
|
||
worker_id=worker_id,
|
||
previous_task_key=task['task_key'],
|
||
status='success'
|
||
)
|
||
stats = self.dispatcher.get_statistics()
|
||
assert stats['completed'] == 1
|
||
assert stats['running'] <= 1 # 可能领了新任务
|
||
|
||
|
||
# ==================== PubSub集成测试 ====================
|
||
@skip_no_redis
|
||
class TestPubSubIntegration:
|
||
"""通过PubSub走完整Master-Worker流程
|
||
|
||
启动 DispatcherService 作为Master，
|
||
多个模拟Worker线程通过PubSub通信
|
||
"""
|
||
|
||
def setup_method(self):
|
||
"""每个测试前清空DB 15"""
|
||
helper.flush()
|
||
self.temp_dir = tempfile.mkdtemp()
|
||
config.FAILED_TASKS_CSV = os.path.join(self.temp_dir, 'failed_tasks.csv')
|
||
config.SUCCESS_TASKS_CSV = os.path.join(self.temp_dir, 'success_tasks.csv')
|
||
|
||
def teardown_method(self):
|
||
helper.flush()
|
||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||
|
||
def _graceful_stop_service(self, service, clear_data=True):
|
||
"""优雅停止Master服务，避免daemon线程I/O关闭警告
|
||
|
||
先设置 running=False 让监听线程自然退出循环（每次 get_message 超时1秒），
|
||
等待线程退出后再关闭连接，避免在阻塞读取时关闭socket。
|
||
"""
|
||
service.running = False
|
||
# 等待监听线程退出（线程轮询间隔为1秒，等2秒足够）
|
||
time.sleep(2)
|
||
# 线程已退出，安全关闭pubsub连接
|
||
with service._lock:
|
||
for pubsub in service.pubsub_list:
|
||
try:
|
||
pubsub.unsubscribe()
|
||
pubsub.close()
|
||
except Exception:
|
||
pass
|
||
service.pubsub_list.clear()
|
||
# 按需清理数据
|
||
if clear_data:
|
||
service._clear_all_data()
|
||
|
||
def _create_worker_redis(self):
|
||
"""创建Worker使用的Redis连接（DB 15）"""
|
||
return redis.Redis(host='localhost', port=6379, db=TEST_REDIS_DB, decode_responses=True)
|
||
|
||
def _simulate_worker(self, worker_idx, redis_conn, results, errors, stop_event):
|
||
"""模拟单个Worker通过PubSub与Master通信
|
||
|
||
Args:
|
||
worker_idx: Worker编号
|
||
redis_conn: Redis连接
|
||
results: 存放结果的字典
|
||
errors: 存放错误的列表
|
||
stop_event: 停止信号
|
||
"""
|
||
worker_id = f"192.168.1.{worker_idx}_AA:BB:CC:DD:EE:{worker_idx:02d}"
|
||
tasks_done = []
|
||
|
||
try:
|
||
# --- init ---
|
||
pubsub = redis_conn.pubsub()
|
||
response_channel = f"worker:init:response:{worker_id}"
|
||
pubsub.subscribe(response_channel)
|
||
|
||
init_request = {
|
||
'worker_id': worker_id,
|
||
'ip_address': f"192.168.1.{worker_idx}",
|
||
'mac_address': f"AA:BB:CC:DD:EE:{worker_idx:02d}",
|
||
'hostname': f"test-pc-{worker_idx}",
|
||
'platform': "Windows"
|
||
}
|
||
redis_conn.publish("worker:init", json.dumps(init_request))
|
||
|
||
# 等待响应
|
||
task = None
|
||
start = time.time()
|
||
while time.time() - start < 10:
|
||
msg = pubsub.get_message(timeout=1)
|
||
if msg and msg['type'] == 'message':
|
||
response = json.loads(msg['data'])
|
||
task = response.get('task')
|
||
break
|
||
|
||
pubsub.unsubscribe()
|
||
pubsub.close()
|
||
|
||
if task is None:
|
||
results[worker_id] = tasks_done
|
||
return
|
||
|
||
# --- report 循环 ---
|
||
while task and not stop_event.is_set():
|
||
tasks_done.append(task['task_key'])
|
||
time.sleep(0.2) # 模拟处理
|
||
|
||
# 发送report
|
||
pubsub2 = redis_conn.pubsub()
|
||
report_response_channel = f"worker:report:response:{worker_id}"
|
||
pubsub2.subscribe(report_response_channel)
|
||
|
||
report_request = {
|
||
'worker_id': worker_id,
|
||
'previous_task_key': task['task_key'],
|
||
'status': 'success',
|
||
'message': ''
|
||
}
|
||
redis_conn.publish("worker:report", json.dumps(report_request))
|
||
|
||
task = None
|
||
start = time.time()
|
||
while time.time() - start < 10:
|
||
msg = pubsub2.get_message(timeout=1)
|
||
if msg and msg['type'] == 'message':
|
||
response = json.loads(msg['data'])
|
||
task = response.get('task')
|
||
break
|
||
|
||
pubsub2.unsubscribe()
|
||
pubsub2.close()
|
||
|
||
results[worker_id] = tasks_done
|
||
|
||
except Exception as e:
|
||
errors.append(f"Worker {worker_idx}: {e}")
|
||
|
||
def test_full_pubsub_flow(self):
|
||
"""完整PubSub流程: Master启动 → 3个Worker通过PubSub领取并完成所有任务"""
|
||
# 启动Master
|
||
service = create_test_service()
|
||
service.dispatcher.load_tasks_from_csv(TEST_CSV_PATH)
|
||
|
||
# 启动监听线程
|
||
threads = service.start_listeners()
|
||
time.sleep(0.5) # 等待监听线程就绪
|
||
|
||
# 启动3个模拟Worker
|
||
results = {}
|
||
errors = []
|
||
stop_event = threading.Event()
|
||
worker_threads = []
|
||
|
||
for i in range(1, 4):
|
||
worker_redis = self._create_worker_redis()
|
||
t = threading.Thread(
|
||
target=self._simulate_worker,
|
||
args=(i, worker_redis, results, errors, stop_event)
|
||
)
|
||
worker_threads.append(t)
|
||
t.start()
|
||
|
||
# 等待所有Worker完成
|
||
for t in worker_threads:
|
||
t.join(timeout=60)
|
||
|
||
# 优雅停止Master（先让线程自然退出再关闭连接）
|
||
self._graceful_stop_service(service, clear_data=False)
|
||
|
||
# 验证
|
||
assert len(errors) == 0, f"不应有错误: {errors}"
|
||
|
||
all_tasks = []
|
||
for tasks in results.values():
|
||
all_tasks.extend(tasks)
|
||
|
||
assert len(all_tasks) == 10, \
|
||
f"所有10个任务应被完成，实际: {len(all_tasks)}，详情: {results}"
|
||
assert len(set(all_tasks)) == 10, "不应有重复分配"
|
||
|
||
def test_pubsub_worker_stop(self):
|
||
"""PubSub流程: Worker领取任务后上报stop"""
|
||
service = create_test_service()
|
||
service.dispatcher.load_tasks_from_csv(TEST_CSV_PATH)
|
||
threads = service.start_listeners()
|
||
time.sleep(0.5)
|
||
|
||
worker_id = "192.168.1.1_AA:BB:CC:DD:EE:01"
|
||
worker_redis = self._create_worker_redis()
|
||
|
||
try:
|
||
# init
|
||
pubsub = worker_redis.pubsub()
|
||
pubsub.subscribe(f"worker:init:response:{worker_id}")
|
||
worker_redis.publish("worker:init", json.dumps({
|
||
'worker_id': worker_id,
|
||
'ip_address': "192.168.1.1",
|
||
'mac_address': "AA:BB:CC:DD:EE:01",
|
||
'hostname': "test-pc-stop",
|
||
'platform': "Windows"
|
||
}))
|
||
|
||
task = None
|
||
start = time.time()
|
||
while time.time() - start < 10:
|
||
msg = pubsub.get_message(timeout=1)
|
||
if msg and msg['type'] == 'message':
|
||
task = json.loads(msg['data']).get('task')
|
||
break
|
||
pubsub.unsubscribe()
|
||
pubsub.close()
|
||
|
||
assert task is not None, "Worker应能领到任务"
|
||
|
||
# report stop
|
||
pubsub2 = worker_redis.pubsub()
|
||
pubsub2.subscribe(f"worker:report:response:{worker_id}")
|
||
worker_redis.publish("worker:report", json.dumps({
|
||
'worker_id': worker_id,
|
||
'previous_task_key': task['task_key'],
|
||
'status': 'stop',
|
||
'message': '设备异常'
|
||
}))
|
||
|
||
response_task = "NOT_RECEIVED"
|
||
start = time.time()
|
||
while time.time() - start < 10:
|
||
msg = pubsub2.get_message(timeout=1)
|
||
if msg and msg['type'] == 'message':
|
||
response_task = json.loads(msg['data']).get('task')
|
||
break
|
||
pubsub2.unsubscribe()
|
||
pubsub2.close()
|
||
|
||
assert response_task is None, "stop后Master不应返回新任务"
|
||
|
||
# 验证Worker已被移除
|
||
time.sleep(0.5)
|
||
workers = service.dispatcher.get_registered_workers()
|
||
worker_ids = [w['worker_id'] for w in workers]
|
||
assert worker_id not in worker_ids, "Worker应被移除"
|
||
|
||
finally:
|
||
self._graceful_stop_service(service, clear_data=True)
|
||
|
||
|
||
# ==================== 入口 ====================
|
||
if __name__ == '__main__':
|
||
# 支持直接运行: python tests/test_master_worker.py
|
||
pytest.main([__file__, '-v', '--tb=short'])
|