#!/usr/bin/env bats
# tests/pcie-probe.bats — black-box tests for scripts/pcie-probe.sh (Refs #53)
#
# Contract: measure host→device bandwidth on an RDNA4 box, and refuse to measure while the GPU is
# in use. Fleet law: never benchmark beside a live service — so the guard runs BEFORE anything is
# built or run, and the tool stops nothing and starts nothing (auto-stopping a production service is
# the thing it must not do unasked).
#
# Exit status: 0 measured · 2 usage · 3 refused (GPU busy) · 4 no toolchain / build failed ·
# 5 a probe ran and failed.
#
# hipcc and the compiled probes are stubbed: this box has no HIP SDK (it is a client laptop) and a
# real measurement belongs on a machine with the cards, beside no running model.

setup() {
  TEST_TMP="$(mktemp -d)"
  SCRIPT="$BATS_TEST_DIRNAME/../scripts/pcie-probe.sh"
  SRC="$TEST_TMP/src"
  BUILD="$TEST_TMP/build"
  SYSFS="$TEST_TMP/drm"
  BIN="$TEST_TMP/bin"
  mkdir -p "$SRC" "$BUILD" "$BIN" "$SYSFS/card0/device" "$SYSFS/card1/device"

  # Fixtures stand in for the real sysfs files the driver reads.
  printf '8\n'            > "$SYSFS/card0/device/current_link_width"
  printf '32.0 GT/s PCIe\n' > "$SYSFS/card0/device/current_link_speed"
  printf '16\n'           > "$SYSFS/card0/device/max_link_width"
  printf '32.0 GT/s PCIe\n' > "$SYSFS/card0/device/max_link_speed"
  printf '0\n'            > "$SYSFS/card0/device/gpu_busy_percent"
  printf '8\n'            > "$SYSFS/card1/device/current_link_width"
  printf '32.0 GT/s PCIe\n' > "$SYSFS/card1/device/current_link_speed"
  printf '16\n'           > "$SYSFS/card1/device/max_link_width"
  printf '32.0 GT/s PCIe\n' > "$SYSFS/card1/device/max_link_speed"
  printf '0\n'            > "$SYSFS/card1/device/gpu_busy_percent"

  # Probe sources: the real .hip files are shipped by the script's own directory, so the fixtures
  # only need to exist for the build step to be invoked.
  : > "$SRC/pcie_probe.hip"; : > "$SRC/pcie_bw.hip"; : > "$SRC/hip_smoke.hip"

  # Stub toolchain: records that it was called, and emits a probe binary that records that IT ran.
  cat > "$BIN/hipcc" <<'STUB'
#!/usr/bin/env bash
[ "${STUB_HIPCC_FAILS:-0}" = 1 ] && { echo "stub: no matching GPU target" >&2; exit 1; }
out=""
while [ $# -gt 0 ]; do [ "$1" = "-o" ] && out="$2"; shift; done
[ -n "$out" ] || exit 2
: > "$PCIE_PROBE_BUILD/.built"
{ echo '#!/usr/bin/env bash'
  echo ': > "$PCIE_PROBE_BUILD/.ran"'
  # HIP rule, modelled because it is the trap: with HIP_VISIBLE_DEVICES set, the ONLY
  # valid ordinal is 0. A driver that masks the device and still passes its real ordinal
  # asks for a card the runtime cannot see — which is why the stub used to report the
  # env var instead of the argument, and could not see the collision.
  echo 'dev="${1:-0}"'
  echo 'if [ -n "${HIP_VISIBLE_DEVICES:-}" ] && [ "$dev" != 0 ]; then'
  echo '  printf "setDevice(%s) failed\n" "$dev" >&2; exit 1'
  echo 'fi'
  echo '[ "${STUB_PROBE_FAILS:-0}" = 1 ] && { echo "stub probe died mid-measurement" >&2; exit 1; }'
  echo 'echo "PROBE RAN dev=$dev bin=$(basename "$0")"'
  echo 'echo "  slab 2.344 MiB: 31.4 GB/s"'
} > "$out"
chmod +x "$out"
STUB
  cat > "$BIN/fuser" <<'STUB'
#!/usr/bin/env bash
[ -n "${STUB_KFD_PIDS:-}" ] && printf '%s\n' $STUB_KFD_PIDS
exit 0
STUB
  chmod +x "$BIN/hipcc" "$BIN/fuser"

  export PATH="$BIN:$PATH" PCIE_PROBE_SRC="$SRC" PCIE_PROBE_BUILD="$BUILD" \
         PCIE_PROBE_SYSFS="$SYSFS" PCIE_PROBE_KFD="$TEST_TMP/kfd" PCIE_PROBE_HIPCC="$BIN/hipcc"
  : > "$TEST_TMP/kfd"
}

teardown() { rm -rf "$TEST_TMP"; }

@test "a missing HIP toolchain is a loud failure, not a silent skip" {
  run env PCIE_PROBE_HIPCC="$TEST_TMP/nope-hipcc" "$SCRIPT"
  [ "$status" -eq 4 ]
  [[ "$output" == *"hip"* ]]
  [ ! -e "$BUILD/.built" ]
}

@test "a failed build is reported as a failure" {
  run env STUB_HIPCC_FAILS=1 "$SCRIPT"
  [ "$status" -eq 4 ]
  [[ "$output" == *"build"* ]]
}

@test "a GPU held by another process refuses the measurement before anything is built" {
  # Order is the contract: build-and-then-ask would already have loaded the card beside a live
  # service. The .built marker proves the guard ran first.
  run env STUB_KFD_PIDS="4242" "$SCRIPT"
  [ "$status" -eq 3 ]
  [[ "$output" == *"busy"* ]]
  [ ! -e "$BUILD/.built" ]
}

@test "a busy GPU percent refuses too" {
  printf '91\n' > "$SYSFS/card0/device/gpu_busy_percent"
  run "$SCRIPT"
  [ "$status" -eq 3 ]
  [[ "$output" == *"card0"* ]]
  [ ! -e "$BUILD/.built" ]
}

@test "--force overrides the guard for a deliberately contended measurement" {
  run env STUB_KFD_PIDS="4242" "$SCRIPT" --force
  [ "$status" -eq 0 ]
  [ -e "$BUILD/.ran" ]
}

@test "an idle GPU builds and runs the probe" {
  run "$SCRIPT"
  [ "$status" -eq 0 ]
  [ -e "$BUILD/.built" ]
  [ -e "$BUILD/.ran" ]
  [[ "$output" == *"PROBE RAN"* ]]
  [[ "$output" == *"31.4 GB/s"* ]]   # probe output passes through untouched
}

@test "--build-only compiles without touching the card" {
  run "$SCRIPT" --build-only
  [ "$status" -eq 0 ]
  [ -e "$BUILD/.built" ]
  [ ! -e "$BUILD/.ran" ]
}

@test "the link topology is reported from sysfs, current against max" {
  run "$SCRIPT"
  [[ "$output" == *"card0"*"x8"*"32.0 GT/s"* ]]
  [[ "$output" == *"max x16"* ]]
}

@test "each requested device is measured separately" {
  # One shared run cannot attribute bandwidth to a card; the whole point is per-card numbers.
  run "$SCRIPT" --device 0 --device 1
  [ "$status" -eq 0 ]
  [[ "$output" == *"dev=0"* ]]
  [[ "$output" == *"dev=1"* ]]
}

@test "a device ordinal the environment mask would hide is still measured" {
  # Reproduction, not preference: HIP_VISIBLE_DEVICES=1 leaves ordinal 0 as the only
  # selectable card, so masking AND passing the ordinal asks hipSetDevice(1) for a
  # device the runtime cannot see. Per-card bandwidth was unobtainable that way.
  run "$SCRIPT" --device 0 --device 1
  [ "$status" -eq 0 ]
  [[ "$output" == *"PROBE RAN dev=1"* ]]
  [[ "$output" != *"setDevice"* ]]
}

@test "a probe that dies is reported as a failed run, not a green one" {
  # The header promises "0 measured". A probe that segfaults (the documented hazard on
  # this ROCm build) must not let the sweep exit 0.
  run env STUB_PROBE_FAILS=1 "$SCRIPT"
  [ "$status" -eq 5 ]
  [[ "$output" == *"failed"* ]]
  [[ "$output" == *"hip_smoke"* ]]
}

@test "the guard refuses when it cannot see who holds the KFD, instead of assuming idle" {
  # Without fuser the /dev/kfd half of the guard cannot run, and gpu_busy_percent reads
  # 0 for a service that is resident but idle — a silent fallback would benchmark beside
  # a live endpoint, which is the one thing this tool's LAW forbids.
  run env PCIE_PROBE_FUSER="$TEST_TMP/nope-fuser" "$SCRIPT"
  [ "$status" -eq 3 ]
  [[ "$output" == *"fuser"* ]]
  [ ! -e "$BUILD/.built" ]
}

@test "no probe sources is a usage-level error, not an empty success" {
  rm -f "$SRC"/*.hip
  run "$SCRIPT"
  [ "$status" -ne 0 ]
  [[ "$output" == *"probe"* ]]
}

@test "an unknown flag is a usage error" {
  run "$SCRIPT" --frobnicate
  [ "$status" -eq 2 ]
}