run_tests.py

import sys, os
sys.path.insert(0, os.path.dirname(__file__))
import numpy as np
import simpy

from schedulers.bandit import PlainUCBScheduler, LinUCBScheduler, DLinUCBScheduler, DKernelUCBScheduler
from schedulers import FIFOScheduler, PriorityScheduler, WFQScheduler
from engine.simulation import run_simulation, run_bandit_simulation, run_warmup_simulation
from engine.access_point import AccessPoint
from wireless_models import PoissonTraffic

results = []

def test(name, fn):
    try:
        fn()
        print(f"  PASS  {name}")
        results.append(True)
    except Exception as e:
        print(f"  FAIL  {name} -- {e}")
        results.append(False)

# ================================================================
print("\n[PlainUCB]")
# ================================================================

def plain_1000():
    p = PlainUCBScheduler(n_users=6)
    ctx = [1.0, 3.0, 0.3, 0.6]
    for _ in range(1000):
        s = p.select_scheduler(ctx)
        p.update(ctx, s, 0.7)
test("1000 updates no crash", plain_1000)

def plain_zero():
    p = PlainUCBScheduler(n_users=6)
    ctx = [0.0, 0.0, 0.0, 0.0]
    for _ in range(10):
        s = p.select_scheduler(ctx)
        p.update(ctx, s, 0.0)
test("zero context + zero reward", plain_zero)

def plain_neg():
    p = PlainUCBScheduler(n_users=6)
    for _ in range(10):
        s = p.select_scheduler([1, 2, 0.3, 0.6])
        p.update([1, 2, 0.3, 0.6], s, -1.0)
test("negative reward -1.0", plain_neg)

def plain_explore():
    p = PlainUCBScheduler(n_users=6)
    seen = set()
    for _ in range(3):
        s = p.select_scheduler([1, 2, 0.3, 0.6])
        seen.add(s.__class__.__name__)
        p.update([1, 2, 0.3, 0.6], s, 0.7)
    assert len(seen) == 3, f"Should explore all 3 arms, got: {seen}"
test("force-explore all 3 arms in first 3 rounds", plain_explore)

# ================================================================
print("\n[LinUCB]")
# ================================================================

def lin_zero():
    l = LinUCBScheduler(n_users=6)
    ctx = [0.0, 0.0, 0.0, 0.0]
    for _ in range(50):
        s = l.select_scheduler(ctx)
        l.update(ctx, s, 0.7)
test("all-zero context 50 updates", lin_zero)

def lin_large():
    l = LinUCBScheduler(n_users=6)
    ctx = [1000.0, 6.0, 1.0, 1.0]
    for _ in range(30):
        s = l.select_scheduler(ctx)
        l.update(ctx, s, 0.8)
test("large context values (1000)", lin_large)

def lin_vary():
    l = LinUCBScheduler(n_users=6)
    for i in range(100):
        ctx = [float(i % 5), float(i % 6 + 1), float(i % 10) * 0.1, 0.5 + float(i % 5) * 0.05]
        s = l.select_scheduler(ctx)
        l.update(ctx, s, 0.6 + (i % 3) * 0.05)
test("varying context + reward 100 steps", lin_vary)

def lin_pd():
    l = LinUCBScheduler(n_users=6)
    ctx = [1.0, 3.0, 0.3, 0.6]
    for _ in range(50):
        s = l.select_scheduler(ctx)
        l.update(ctx, s, 0.7)
    assert np.linalg.det(l._A[0]) > 0
test("A matrix stays positive definite", lin_pd)

# ================================================================
print("\n[DLinUCB]")
# ================================================================

def dlin_200():
    d = DLinUCBScheduler(n_users=6)
    ctx = [1.0, 3.0, 0.3, 0.6]
    for _ in range(200):
        s = d.select_scheduler(ctx)
        d.update(ctx, s, 0.7)
    assert np.linalg.det(d._A[0] + 1e-6 * np.eye(4)) > 0
test("200 updates A valid (gamma=0.95)", dlin_200)

def dlin_gamma01():
    d = DLinUCBScheduler(n_users=6, gamma=0.1)
    for _ in range(100):
        s = d.select_scheduler([1, 2, 0.3, 0.6])
        d.update([1, 2, 0.3, 0.6], s, 0.7)
test("aggressive gamma=0.1 no singular matrix", dlin_gamma01)

def dlin_gamma05():
    d = DLinUCBScheduler(n_users=6, gamma=0.5)
    for _ in range(100):
        s = d.select_scheduler([1, 2, 0.3, 0.6])
        d.update([1, 2, 0.3, 0.6], s, 0.7)
test("medium gamma=0.5 no crash", dlin_gamma05)

def dlin_nonstat():
    d = DLinUCBScheduler(n_users=6)
    for _ in range(30):
        ctx = [2.0, 5.0, 0.4, 0.6]
        s = d.select_scheduler(ctx)
        reward = 0.9 if s.__class__.__name__ == "WFQScheduler" else 0.4
        d.update(ctx, s, reward)
    for _ in range(30):
        ctx = [1.0, 3.0, 0.2, 0.7]
        s = d.select_scheduler(ctx)
        reward = 0.9 if s.__class__.__name__ == "FIFOScheduler" else 0.4
        d.update(ctx, s, reward)
test("adapts after reward signal flip (non-stationary)", dlin_nonstat)

# ================================================================
print("\n[DKernelUCB]")
# ================================================================

def dker_cold():
    dk = DKernelUCBScheduler(n_users=6)
    s = dk.select_scheduler([1, 3, 0.3, 0.6])
    assert s is not None
test("cold start select no crash", dker_cold)

def dker_empty():
    dk = DKernelUCBScheduler(n_users=6)
    dk.pretrain([])
    s = dk.select_scheduler([1, 3, 0.3, 0.6])
    assert s is not None
test("empty pretrain list", dker_empty)

def dker_unknown():
    dk = DKernelUCBScheduler(n_users=6)
    dk.pretrain([{"context": [1,2,0.3,0.6], "scheduler": "UnknownXYZ", "reward": 0.7}])
    s = dk.select_scheduler([1, 3, 0.3, 0.6])
    assert s is not None
    assert all(len(dk._data[i]) == 0 for i in range(3)), "Unknown scheduler should be skipped"
test("unknown scheduler in pretrain skipped", dker_unknown)

def dker_100():
    dk = DKernelUCBScheduler(n_users=6)
    for i in range(100):
        ctx = [float(i % 5), 3.0, 0.3, 0.6]
        s = dk.select_scheduler(ctx)
        dk.update(ctx, s, 0.7)
test("100 updates growing kernel matrix", dker_100)

def dker_sigma():
    dk = DKernelUCBScheduler(n_users=6)
    ctx = [1.0, 3.0, 0.3, 0.6]
    for _ in range(20):
        s = dk.select_scheduler(ctx)
        dk.update(ctx, s, 0.7)
    assert dk.sigma > 0, f"sigma={dk.sigma}"
test("identical contexts sigma stays > 0", dker_sigma)

def dker_zero_reward():
    dk = DKernelUCBScheduler(n_users=6)
    for i in range(20):
        s = dk.select_scheduler([1, 3, 0.3, float(i) * 0.01])
        dk.update([1, 3, 0.3, float(i) * 0.01], s, 0.0)
test("all zero rewards no crash", dker_zero_reward)

def dker_pretrain_learns():
    dk = DKernelUCBScheduler(n_users=6)
    obs = []
    for i in range(8):
        obs.append({"context": [1.5, 3.0, 0.3, 0.60+i*0.001], "scheduler": "FIFOScheduler",     "reward": 0.68})
        obs.append({"context": [1.5, 3.0, 0.3, 0.61+i*0.001], "scheduler": "PriorityScheduler", "reward": 0.70})
        obs.append({"context": [1.5, 3.0, 0.3, 0.62+i*0.001], "scheduler": "WFQScheduler",      "reward": 0.73})
    dk.pretrain(obs)
    assert all(len(dk._data[i]) == 8 for i in range(3)), f"Expected 8 obs per arm, got {[len(dk._data[i]) for i in range(3)]}"
    dk._tune_sigma()
    x = np.array([1.5, 3.0, 0.3, 0.62])
    mu_fifo,  _ = dk._predict(0, x)
    mu_wfq,   _ = dk._predict(2, x)
    assert mu_wfq > mu_fifo, f"WFQ mu ({mu_wfq:.4f}) should be > FIFO mu ({mu_fifo:.4f})"
test("pretrain loads data correctly and WFQ has highest mean reward", dker_pretrain_learns)

def dker_no_autotune():
    dk = DKernelUCBScheduler(n_users=6, auto_tune_sigma=False, sigma=0.5)
    for _ in range(20):
        s = dk.select_scheduler([1, 3, 0.3, 0.6])
        dk.update([1, 3, 0.3, 0.6], s, 0.7)
    assert dk.sigma == 0.5, "sigma should not change"
test("auto_tune_sigma=False sigma stays fixed", dker_no_autotune)

def dker_discount_weights():
    dk = DKernelUCBScheduler(n_users=6, gamma=0.9)
    ctx = [1.0, 3.0, 0.3, 0.6]
    # Force same arm (arm 0 = FIFO) for both updates to verify discount
    arm0_sched = dk.schedulers[0]
    dk.update(ctx, arm0_sched, 0.7)         # arm 0: weight = 1.0
    w_before = dk._data[0][0][2]
    dk.update(ctx, arm0_sched, 0.8)         # arm 0: existing weight *= 0.9
    w_after = dk._data[0][0][2]
    assert w_after < w_before, f"Weight should decay: before={w_before}, after={w_after}"
    assert abs(w_after - w_before * 0.9) < 1e-9, f"Weight should decay by gamma=0.9"
test("discount weights decay correctly with gamma", dker_discount_weights)

# ================================================================
print("\n[Simulation Stress]")
# ================================================================

def sim_single():
    env = simpy.Environment()
    ap = AccessPoint(env, n_users=1, link_capacity_bps=1_000_000, scheduler=FIFOScheduler(n_users=1))
    env.process(PoissonTraffic(env, ap, user_id=0, rate_pps=100, packet_size_bits=8000).run())
    env.process(ap.service_process())
    env.run(until=10)
    assert ap.fairness_evaluation()["jains_fairness_index"] == 1.0
test("single user JFI=1.0", sim_single)

def sim_congestion():
    r = run_simulation(FIFOScheduler(n_users=6), sim_time=10, link_capacity_bps=100_000)
    assert r["normalized_throughput"] <= 1.0
test("extreme congestion 0.1Mbps", sim_congestion)

def sim_fast():
    r = run_simulation(WFQScheduler(n_users=6), sim_time=10, link_capacity_bps=100_000_000)
    assert r["normalized_throughput"] < 1.0
test("100Mbps link no crash", sim_fast)

def sim_rapid():
    dk = DKernelUCBScheduler(n_users=6)
    r = run_bandit_simulation(dk, dk.schedulers, sim_time=20, epoch_interval=0.5)
    assert "composite_reward" in r
test("rapid 0.5s epoch switching", sim_rapid)

def sim_warmup_empty():
    obs = run_warmup_simulation([FIFOScheduler(6)], warmup_time=3.0, epoch_interval=5.0)
    assert isinstance(obs, list)
test("warmup shorter than epoch (empty obs)", sim_warmup_empty)

def sim_fw0():
    r = run_simulation(WFQScheduler(n_users=6), sim_time=10, fairness_weight=0.0)
    assert abs(r["composite_reward"] - r["normalized_throughput"]) < 1e-10
test("fairness_weight=0.0 composite==throughput", sim_fw0)

def sim_fw1():
    r = run_simulation(WFQScheduler(n_users=6), sim_time=10, fairness_weight=1.0)
    assert abs(r["composite_reward"] - r["jains_fairness_index"]) < 1e-10
test("fairness_weight=1.0 composite==JFI", sim_fw1)

def sim_repro():
    r1 = run_simulation(WFQScheduler(n_users=6), sim_time=10, seed=99)
    r2 = run_simulation(WFQScheduler(n_users=6), sim_time=10, seed=99)
    assert r1["composite_reward"] == r2["composite_reward"]
test("same seed reproducible results", sim_repro)

def sim_diff_seed():
    r1 = run_simulation(WFQScheduler(n_users=6), sim_time=10, seed=1)
    r2 = run_simulation(WFQScheduler(n_users=6), sim_time=10, seed=2)
    assert r1["composite_reward"] != r2["composite_reward"]
test("different seeds give different results", sim_diff_seed)

def sim_swap():
    env = simpy.Environment()
    ap = AccessPoint(env, n_users=6, link_capacity_bps=10_000_000, scheduler=FIFOScheduler(n_users=6))
    ap.swap_scheduler(WFQScheduler(n_users=6))
    assert ap.scheduler.__class__.__name__ == "WFQScheduler"
test("swap_scheduler changes correctly", sim_swap)

def sim_ctx():
    env = simpy.Environment()
    ap = AccessPoint(env, n_users=6, link_capacity_bps=10_000_000, scheduler=FIFOScheduler(n_users=6))
    ctx = ap.get_context()
    assert len(ctx) == 4
    assert all(isinstance(v, float) for v in ctx)
test("get_context returns 4 floats", sim_ctx)

def sim_queue_full():
    env = simpy.Environment()
    ap = AccessPoint(env, n_users=6, link_capacity_bps=10_000_000,
                     scheduler=FIFOScheduler(n_users=6), max_queue_size=1)
    env.process(PoissonTraffic(env, ap, user_id=0, rate_pps=500, packet_size_bits=8000).run())
    env.process(ap.service_process())
    env.run(until=5)
    assert ap.stats[0].dropped_packets > 0, "Queue full should drop packets"
test("queue overflow drops packets correctly", sim_queue_full)

# ================================================================
print()
total = len(results)
passed = sum(results)
print("=" * 50)
print(f"TOTAL: {passed}/{total} passed  |  {total - passed} failed")
if total - passed == 0:
    print("ALL TESTS PASSED")