test_benchmark_config.py

#!/usr/bin/env python3
"""
Тестовый скрипт для проверки конфигурации бенчмарка

Этот скрипт проверяет, что обновленная конфигурация бенчмарка
корректно загружается и используется существующим кодом.
"""

import sys
import logging
import argparse
import random
import time
from pprint import pprint

# Импорт конфигурации
from config import benchmark_config, cloud_configs
from incremental_reporting import (
    get_benchmark_files,
    save_benchmark_result,
    update_benchmark_state,
    is_benchmark_completed,
    load_results
)

# Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler()
    ]
)
logger = logging.getLogger('pg_benchmark_test')

def validate_config():
    """Проверка наличия всех необходимых параметров в конфигурации."""
    required_params = [
        "scale_factor", "client_counts", "thread_counts",
        "transactions_per_client", "run_time", "use_run_time",
        "iterations", "initialization", "vacuum_before_tests",
        "drop_tables_after_tests", "workload_types",
        "connection_timeout", "connection_retries", "use_ssl",
        "progress_reporting", "progress_interval",
        "test_network_latency", "ping_count"
    ]

    missing_params = [param for param in required_params if param not in benchmark_config]

    if missing_params:
        logger.error(f"Отсутствуют обязательные параметры: {', '.join(missing_params)}")
        return False

    logger.info("Все обязательные параметры присутствуют в конфигурации")
    return True

def print_config_summary():
    """Вывод сводки по конфигурации бенчмарка."""
    print("\n=== Сводка по конфигурации бенчмарка ===\n")

    # Базовые параметры
    print(f"Размер базы данных: scale_factor = {benchmark_config['scale_factor']}")
    print(f"Клиенты: {benchmark_config['client_counts']}")
    print(f"Потоки: {benchmark_config['thread_counts']}")

    # Продолжительность
    if benchmark_config['use_run_time']:
        print(f"Продолжительность: {benchmark_config['run_time']} секунд")
    else:
        print(f"Транзакций на клиента: {benchmark_config['transactions_per_client']}")

    # Итерации
    print(f"Итерации: {benchmark_config['iterations']}")
    if 'warmup_iterations' in benchmark_config:
        print(f"Разогревочные итерации: {benchmark_config['warmup_iterations']}")

    # Рабочие нагрузки
    print("\nРабочие нагрузки:")
    for workload in benchmark_config['workload_types']:
        desc = workload.get('description', 'Нет описания')
        print(f"  - {workload['name']}: {desc}")

    # Подготовка БД
    print("\nПодготовка базы данных:")
    print(f"  - Инициализация: {benchmark_config['initialization']}")
    print(f"  - Vacuum перед тестами: {benchmark_config['vacuum_before_tests']}")
    if 'analyze_before_tests' in benchmark_config:
        print(f"  - Analyze перед тестами: {benchmark_config['analyze_before_tests']}")
    print(f"  - Удаление таблиц после тестов: {benchmark_config['drop_tables_after_tests']}")

    # Настройки соединения
    print("\nНастройки соединения:")
    print(f"  - Таймаут соединения: {benchmark_config['connection_timeout']} сек")
    print(f"  - Попытки соединения: {benchmark_config['connection_retries']}")
    print(f"  - Использовать SSL: {benchmark_config['use_ssl']}")
    print(f"  - Режим SSL: {benchmark_config['ssl_mode']}")

    # Расширенные параметры
    print("\nРасширенные параметры:")
    for key in benchmark_config:
        if key not in [
            'scale_factor', 'client_counts', 'thread_counts',
            'transactions_per_client', 'run_time', 'use_run_time',
            'iterations', 'initialization', 'vacuum_before_tests',
            'drop_tables_after_tests', 'workload_types',
            'connection_timeout', 'connection_retries', 'use_ssl',
            'progress_reporting', 'progress_interval',
            'test_network_latency', 'ping_count', 'warmup_iterations',
            'analyze_before_tests', 'scale_factors', 'run_time_options',
            'initialization_options', 'ssl_mode'
        ]:
            print(f"  - {key}: {type(benchmark_config[key]).__name__}")

    # Провайдеры
    print("\nНастроенные провайдеры PostgreSQL:")
    for config in cloud_configs:
        if isinstance(config, dict) and 'name' in config:
            provider_info = f"  - {config['name']} ({config['instance_type']})"

            # Добавление информации о SSL, если настроено
            ssl_info = []
            if 'ssl_mode' in config:
                ssl_info.append(f"SSL: {config['ssl_mode']}")
            if 'ssl_root_cert' in config:
                ssl_info.append(f"Cert: {config['ssl_root_cert']}")

            if ssl_info:
                provider_info += f" [{', '.join(ssl_info)}]"

            print(provider_info)

    print("\n=== Конец сводки ===\n")

def simulate_benchmark(
    config,
    client_count,
    thread_count,
    workload_name,
    iteration,
    results_file,
    state_file,
    delay=2,
    fail_probability=0.0
):
    """
    Симуляция выполнения бенчмарка с сохранением результатов.

    Args:
        config: Конфигурация провайдера
        client_count: Количество клиентов
        thread_count: Количество потоков
        workload_name: Имя рабочей нагрузки
        iteration: Номер итерации
        results_file: Имя файла результатов
        state_file: Имя файла состояния
        delay: Задержка в секундах для симуляции длительного бенчмарка
        fail_probability: Вероятность сбоя (0.0 - 1.0)
    """
    # Проверка, был ли уже выполнен этот бенчмарк
    if is_benchmark_completed(
        config["name"],
        client_count,
        thread_count,
        workload_name,
        iteration,
        state_file
    ):
        logger.info(
            f"Пропуск уже выполненного бенчмарка: {config['name']}, "
            f"клиенты={client_count}, потоки={thread_count}, "
            f"нагрузка={workload_name}, итерация={iteration}"
        )
        return

    logger.info(
        f"Запуск бенчмарка: {config['name']}, "
        f"клиенты={client_count}, потоки={thread_count}, "
        f"нагрузка={workload_name}, итерация={iteration}"
    )

    # Симуляция случайного сбоя
    if random.random() < fail_probability:
        logger.error(f"Симуляция сбоя для бенчмарка!")
        raise Exception("Симулированный сбой бенчмарка")

    # Симуляция выполнения бенчмарка
    time.sleep(delay)

    # Создание симулированного результата
    tps = random.uniform(100, 500)
    latency = random.uniform(10, 50)

    result = {
        "provider": config["name"],
        "instance_type": config["instance_type"],
        "version": config["version"],
        "client_count": client_count,
        "thread_count": thread_count,
        "workload": workload_name,
        "iteration": iteration,
        "tps": tps,
        "latency_avg": latency,
        "latency_stddev": latency * 0.1,
        "cost_per_hour": config["rub_per_hour"],
        "cost_performance": tps / (config["rub_per_hour"] / 3600),
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "duration": delay,
        "simulated": True
    }

    # Сохранение результата и обновление состояния
    save_benchmark_result(result, results_file)
    update_benchmark_state(
        config["name"],
        client_count,
        thread_count,
        workload_name,
        iteration,
        state_file
    )

    logger.info(
        f"Бенчмарк завершен: {config['name']}, "
        f"клиенты={client_count}, потоки={thread_count}, "
        f"нагрузка={workload_name}, итерация={iteration}, "
        f"TPS: {tps:.2f}, Задержка: {latency:.2f} мс"
    )

def run_simulation(configs, max_benchmarks=10, fail_after=None, delay=2, fail_probability=0.0):
    """
    Запуск симуляции бенчмаркинга с инкрементальным сохранением.

    Args:
        configs: Список конфигураций провайдеров
        max_benchmarks: Максимальное количество бенчмарков для симуляции
        fail_after: Номер бенчмарка, после которого симулировать сбой
        delay: Задержка в секундах для симуляции длительного бенчмарка
        fail_probability: Вероятность сбоя для каждого бенчмарка
    """
    # Получение файлов для сохранения результатов и состояния
    results_file, state_file = get_benchmark_files()

    # Загрузка существующих результатов
    existing_results = load_results(results_file)
    logger.info(f"Загружено {len(existing_results)} существующих результатов")

    # Счетчик выполненных бенчмарков
    benchmark_count = 0

    try:
        # Перебор всех комбинаций параметров
        for config in configs:
            for client_count in [2, 4, 8]:
                for thread_count in [1, 2]:
                    for workload_name in ["default", "simple_update", "select_only"]:
                        for iteration in range(1, 3):
                            # Симуляция бенчмарка
                            simulate_benchmark(
                                config,
                                client_count,
                                thread_count,
                                workload_name,
                                iteration,
                                results_file,
                                state_file,
                                delay=delay,
                                fail_probability=fail_probability
                            )

                            benchmark_count += 1

                            # Проверка на достижение максимального количества бенчмарков
                            if max_benchmarks and benchmark_count >= max_benchmarks:
                                logger.info(f"Достигнуто максимальное количество бенчмарков ({max_benchmarks})")
                                return

                            # Симуляция сбоя после указанного количества бенчмарков
                            if fail_after and benchmark_count >= fail_after:
                                logger.warning(f"Симуляция сбоя после {fail_after} бенчмарков")
                                raise KeyboardInterrupt("Симулированное прерывание")

    except KeyboardInterrupt:
        logger.warning("Бенчмаркинг прерван пользователем или симулированным сбоем")

    finally:
        # Загрузка и вывод итоговых результатов
        final_results = load_results(results_file)
        logger.info(f"Всего сохранено {len(final_results)} результатов в {results_file}")
        logger.info(f"Файл состояния: {state_file}")

def simulate_benchmark_run(quick_mode=False):
    """
    Симуляция запуска бенчмарка для проверки конфигурации.

    Args:
        quick_mode: Если True, использует минимальные параметры для быстрой проверки
    """
    # Создание копии конфигурации для симуляции
    test_config = benchmark_config.copy()

    if quick_mode:
        # Установка минимальных параметров для быстрой проверки
        test_config['client_counts'] = [2]
        test_config['thread_counts'] = [2]
        test_config['iterations'] = 1
        test_config['run_time'] = 5
        test_config['workload_types'] = [wl for wl in test_config['workload_types'] if wl['name'] == 'default']

    print("\n=== Симуляция запуска бенчмарка ===\n")

    # Вывод параметров симуляции
    print(f"Размер БД: scale_factor = {test_config['scale_factor']}")
    print(f"Клиенты: {test_config['client_counts']}")
    print(f"Потоки: {test_config['thread_counts']}")
    print(f"Итерации: {test_config['iterations']}")

    # Симуляция выполнения для каждого провайдера
    for config in cloud_configs:
        if not isinstance(config, dict) or 'name' not in config:
            continue

        print(f"\nПровайдер: {config['name']}")

        # Симуляция подготовки базы данных
        if test_config['initialization']:
            print(f"  Инициализация таблиц pgbench (scale_factor={test_config['scale_factor']})")

        if test_config['vacuum_before_tests']:
            print(f"  Выполнение vacuum")

        if test_config.get('analyze_before_tests'):
            print(f"  Выполнение analyze")

        # Симуляция запуска бенчмарков
        for client_count in test_config['client_counts']:
            for thread_count in test_config['thread_counts']:
                for workload in test_config['workload_types']:
                    for iteration in range(1, test_config['iterations'] + 1):
                        workload_name = workload['name']
                        print(f"  Запуск: клиенты={client_count}, потоки={thread_count}, "
                              f"нагрузка={workload_name}, итерация={iteration}")

        # Симуляция очистки
        if test_config['drop_tables_after_tests']:
            print(f"  Удаление таблиц pgbench")

    print("\n=== Конец симуляции ===\n")

def main():
    """Основная функция."""
    parser = argparse.ArgumentParser(description='Тестирование конфигурации бенчмарка PostgreSQL')
    parser.add_argument('--validate', action='store_true', help='Проверить конфигурацию')
    parser.add_argument('--summary', action='store_true', help='Вывести сводку по конфигурации')
    parser.add_argument('--simulate', action='store_true', help='Симулировать запуск бенчмарка')
    parser.add_argument('--quick', action='store_true', help='Использовать минимальные параметры для быстрой проверки')
    parser.add_argument('--max', type=int, default=10, help='Максимальное количество бенчмарков')
    parser.add_argument('--fail-after', type=int, help='Симулировать сбой после указанного количества бенчмарков')
    parser.add_argument('--delay', type=float, default=2.0, help='Задержка в секундах для симуляции бенчмарка')
    parser.add_argument('--fail-probability', type=float, default=0.0, help='Вероятность сбоя для каждого бенчмарка')

    args = parser.parse_args()

    # Если аргументы не указаны, выполнить все проверки
    if not (args.validate or args.summary or args.simulate):
        args.validate = True
        args.summary = True
        args.simulate = True

    if args.validate:
        if not validate_config():
            sys.exit(1)

    if args.summary:
        print_config_summary()

    if args.simulate:
        if args.max or args.fail_after or args.delay or args.fail_probability:
            print("\n=== Тестирование инкрементального бенчмаркинга ===\n")
            print(f"Максимальное количество бенчмарков: {args.max}")
            if args.fail_after:
                print(f"Симуляция сбоя после: {args.fail_after} бенчмарков")
            print(f"Задержка симуляции: {args.delay} сек")
            print(f"Вероятность сбоя: {args.fail_probability}")
            print("\nДля прерывания нажмите Ctrl+C\n")

            # Запуск симуляции
            run_simulation(
                cloud_configs,
                max_benchmarks=args.max,
                fail_after=args.fail_after,
                delay=args.delay,
                fail_probability=args.fail_probability
            )

            print("\n=== Тестирование завершено ===\n")
            print("Для запуска реального бенчмарка используйте: python main.py")
            print("Прерванный бенчмарк будет автоматически продолжен с места остановки")
        else:
            simulate_benchmark_run(args.quick)

    print("Тестирование конфигурации завершено успешно")

if __name__ == "__main__":
    main()