tests/load_test.py - master - NeonBite/big-data-5

tests/load_test.py
87 строк3 KB
NeonBite
feat: report done pushed
20 дек 2025, 17:27
20 дек 2025, 17:270df5e56
import datetime as dt
import logging
import os
import random
import subprocess
import time

from kafka import KafkaProducer


logger = logging.getLogger(__name__)


def setup_logging():
    log_dir = os.environ.get("LOG_DIR", ".logs")
    os.makedirs(log_dir, exist_ok=True)
    script_name = os.path.splitext(os.path.basename(__file__))[0]
    timestamp = dt.datetime.utcnow().strftime("%Y%m%d-%H%M%S")
    log_path = os.path.join(log_dir, f"{script_name}-{timestamp}.log")
    logging.basicConfig(
        level=os.environ.get("LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(name)s - %(message)s",
        handlers=[
            logging.FileHandler(log_path),
            logging.StreamHandler(),
        ],
    )
    logger.info("Logging to %s", log_path)


def send_load(rate: float, duration_sec: int, bootstrap: str, topic: str, prefix: str):
    producer = KafkaProducer(bootstrap_servers=bootstrap)
    start = time.time()
    next_send = start
    sent = 0

    while time.time() - start < duration_sec:
        now = time.time()
        if rate > 0 and now < next_send:
            time.sleep(max(0.0, next_send - now))
            continue
        user_id = f"{prefix}-u{random.randint(1, 100000)}"
        item_id = f"item-{random.randint(1, 100000)}"
        rating = 5.0 if sent % 2 == 0 else 3.0
        ts_sec = int(time.time())
        line = f"{user_id},{item_id},{rating:.1f},{ts_sec}"
        producer.send(topic, line.encode("utf-8"))
        sent += 1
        if rate > 0:
            next_send += 1.0 / rate

    producer.flush()
    elapsed = time.time() - start
    logger.info(
        "Load test sent=%s duration_sec=%.2f rate_target=%.1f rate_actual=%.1f",
        sent,
        elapsed,
        rate,
        sent / elapsed if elapsed > 0 else 0.0,
    )


def main():
    setup_logging()
    bootstrap = os.environ.get("KAFKA_BOOTSTRAP", "localhost:9092")
    topic = os.environ.get("RATINGS_TOPIC", "ratings")
    normal_rate = float(os.environ.get("NORMAL_RATE", "500"))
    peak_rate = float(os.environ.get("PEAK_RATE", str(normal_rate * 2)))
    duration_sec = int(os.environ.get("LOAD_DURATION_SEC", "30"))
    fail_taskmanager = os.environ.get("FAIL_TASKMANAGER", "0") == "1"

    logger.info("Normal load test")
    send_load(normal_rate, duration_sec, bootstrap, topic, "normal")

    logger.info("Peak load test")
    send_load(peak_rate, duration_sec, bootstrap, topic, "peak")

    logger.info("Failure test (manual stop TaskManager or set FAIL_TASKMANAGER=1)")
    if fail_taskmanager:
        subprocess.run(["docker", "compose", "stop", "taskmanager"], check=False)
    send_load(normal_rate, duration_sec, bootstrap, topic, "failure")
    if fail_taskmanager:
        subprocess.run(["docker", "compose", "start", "taskmanager"], check=False)


if __name__ == "__main__":
    main()