#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Olympian LED
"""Tests for rate_kit.py against a fake SCPI server (a threaded socket that behaves like the units we measured).

    python3 -I -m pytest test_rate_kit.py        or        python3 test_rate_kit.py
"""

import json
import os
import re
import shutil
import socket
import sys
import tempfile
import threading
import time
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import rate_kit as rk  # noqa: E402


class FakeClock:
    """Sleeps advance a counter instead of the wall; the fake scope reads the same clock."""

    def __init__(self, start=1792000000.0):
        self.t = 0.0
        self.start = start
        self.lock = threading.Lock()
        self.on_sleep = None

    def now(self):
        with self.lock:
            self.t += 1e-6
            return self.t

    def wall(self):
        return self.start + self.now()

    def sleep(self, s):
        if self.on_sleep:
            self.on_sleep(s)
        time.sleep(0.004)  # lets the fake scope's thread take what was sent before the time moves on
        with self.lock:
            self.t += max(0.0, s)


def short_form(word):
    """SCPI's short form of a mixed-case word (NORMal -> NORM); an all-caps word stays."""
    if word.upper() == word:
        return word
    return "".join(c for c in word if c.isupper() or c.isdigit())


BOOL_KEYS = (":DISPlay", ":STATe", ":ENABle", ":COLor")


class FakeScope:
    """Enough of a DHO / MHO for the kit: settings with Rigol-style answers, run state, a counter that totals AUX OUT
    pulses at `rate` while counting on the AUX input, the edge level per source, an error queue, no answer to an
    unknown query. firmware: 'stock' | 'mod' | 'novaos'."""

    def __init__(self, clock, firmware="stock", model="MHO934", rate=7345.0, aux=4, serial="MHO9X000T00001",
                 seed=None, hide_ignores_scale=True):
        self.clock, self.firmware, self.model, self.rate, self.aux = clock, firmware, model, rate, aux
        self.serial = serial
        self.hide_ignores_scale = hide_ignores_scale
        self.received = []
        self.errors = []
        self.hung = False
        self.lenient = False
        self.setups = {}
        self.hang_after = None
        self.never_answer = set()
        self.lock = threading.Lock()
        maker = "NK TECHNOLOGIES" if firmware == "mod" else "RIGOL TECHNOLOGIES"
        fw = "00.01.03" if firmware == "mod" else "00.01.00"
        self.idn = "%s,%s,%s,%s" % (maker, model, serial, fw)
        n = int(model[5]) if model[:3] in ("DHO", "MHO") and model[5].isdigit() else 4
        self.n = n
        st = {":TRIGger:SWEep": "AUTO", ":TIMebase:MAIN:SCALe": "5.000000E-7", ":TIMebase:MAIN:OFFSet": "0.000000",
              ":ACQuire:MDEPth": "1.0000E+04", ":ACQuire:TYPE": "NORM", ":TRIGger:MODE": "EDGE",
              ":TRIGger:EDGE:SOURce": "CHAN1", ":TRIGger:EDGE:SLOPe": "POS", ":SYSTem:AOUTput": "PFA",
              ":DISPlay:TYPE": "VECT", ":DISPlay:GRADing:TIME": "MIN", ":DISPlay:COLor": "0",
              ":COUNter:ENABle": "0", ":COUNter:SOURce": "CHAN1", ":COUNter:MODE": "FREQ", ":COUNter:TOTalize:ENABle": "0"}
        for i in range(1, n + 1):
            st.update({":CHANnel%d:DISPlay" % i: "1" if i == 1 else "0", ":CHANnel%d:SCALe" % i: "1.000000E+00",
                       ":CHANnel%d:OFFSet" % i: "0.000000E+00", ":CHANnel%d:PROBe" % i: "1.000000E0",
                       ":CHANnel%d:COUPling" % i: "DC", ":CHANnel%d:BWLimit" % i: "OFF"})
        self.gens = ["SOURce1", "SOURce2"] if model.startswith("MHO") else (["SOURce"] if model.endswith("S") else [])
        for g in self.gens:
            st.update({":%s:FUNCtion" % g: "SIN", ":%s:FREQuency" % g: "1.000000E+6", ":%s:VOLTage:AMPLitude" % g: "1.000000E0",
                       ":%s:VOLTage:OFFSet" % g: "0.000000", ":%s:OUTPut:STATe" % g: "0"})
        if seed:
            st.update(seed)
        self.factory = dict(st)
        self.factory.update({":TRIGger:SWEep": "AUTO", ":CHANnel1:SCALe": "5.000000E-02"})
        self.state = st
        self.levels = {i: "0.000000" for i in range(1, n + 1)}
        self.running = False
        self.clear_t = 0.0
        self.srv = socket.socket()
        self.srv.bind(("127.0.0.1", 0))
        self.srv.listen(8)
        self.port = self.srv.getsockname()[1]
        threading.Thread(target=self._accept, daemon=True).start()

    @property
    def addr(self):
        return "127.0.0.1:%d" % self.port

    def close(self):
        try:
            self.srv.close()
        except OSError:
            pass

    def _accept(self):
        while True:
            try:
                c, _ = self.srv.accept()
            except OSError:
                return
            threading.Thread(target=self._serve, args=(c,), daemon=True).start()

    def _serve(self, c):
        f = c.makefile("rb")
        try:
            for raw in f:
                line = raw.decode().strip()
                if not line:
                    continue
                with self.lock:
                    self.received.append(line)
                    if self.hung:
                        continue
                    r = self.handle(line)
                    if self.hang_after and line.startswith(self.hang_after):
                        self.hung = True
                if r is not None:
                    c.sendall((r + "\n").encode())
        except OSError:
            pass
        finally:
            c.close()

    # -- behaviour
    def shown(self, i):
        return self.state.get(":CHANnel%d:DISPlay" % i) == "1"

    def count(self):
        st = self.state
        if not (st[":COUNter:ENABle"] == "1" and st[":COUNter:MODE"] == "TOT" and st[":COUNter:TOTalize:ENABle"] == "1"):
            return 0
        if st[":COUNter:SOURce"] != "CHAN%d" % self.aux or not self.running:
            return 0
        if float(self.levels[self.aux]) < 0.5:
            return 0  # the threshold is the AUX input's trigger level: at 0 V nothing crosses it
        return int(self.rate * (self.clock.now() - self.clear_t))

    def srate(self):
        on = sum(1 for i in range(1, self.n + 1) if self.shown(i))
        if self.state[":COUNter:ENABle"] == "1" and not self.shown(self.aux):
            on += 1  # the counter's channel is digitised
        return "%.6E" % (4e9 if on <= 1 else 2e9 if on == 2 else 1e9)

    def norm_value(self, key, v):
        if key.endswith(BOOL_KEYS):
            return "1" if v.upper() in ("ON", "1") else "0"
        if key.endswith(":MDEPth"):
            m = re.match(r"^(\d+)([kM]?)$", v)
            if m:
                return "%.4E" % (int(m.group(1)) * {"": 1, "k": 1000, "M": 1000000}[m.group(2)])
            return v.upper()
        if re.match(r"^CHAN(nel)?\d$", v, re.I):
            return "CHAN" + v[-1]
        try:
            return "%.6E" % float(v)
        except ValueError:
            return short_form(v)

    def handle(self, line):
        if line.endswith("?") or "? " in line:
            q = line.split("?")[0]
            arg = line.split("?", 1)[1].strip()
            if q in self.never_answer:
                return None
            if line == "*IDN?":
                return self.idn
            if line == "*OPC?":
                return "1"
            if q == ":NOVAos:VERSion":
                return "r81,1.4.5,debug,2026-10-14" if self.firmware == "novaos" else self.unknown()
            if q == ":NOVAos:DENSity:RATE":
                return "%.1f,%.1f,100.00" % (self.rate, self.rate) if self.firmware == "novaos" else self.unknown()
            if q == ":SYSTem:ERRor":
                return self.errors.pop(0) if self.errors else '0,"No error"'
            if q == ":TRIGger:STATus":
                return "TD" if self.running else "STOP"
            if q == ":ACQuire:SRATe":
                return self.srate()
            if q == ":COUNter:CURRent":
                return str(self.count())
            if q == ":MEASure:ITEM":
                return "3.0941E+00"
            if q == ":SYSTem:SETup":
                tok = "SETUP%04d" % len(self.setups)
                self.setups[tok] = (dict(self.state), dict(self.levels), self.running)
                return "#19" + tok
            if q == ":TRIGger:EDGE:LEVel":
                return self.levels[int(self.state[":TRIGger:EDGE:SOURce"][-1])]
            if q in self.state:
                return self.state[q]
            if self.lenient:
                return "0"
            return self.unknown()
        if line.upper().startswith((":DISP:DATA", ":DISPLAY:DATA")):
            raise AssertionError("the kit sent a screenshot query")
        k, _, v = line.partition(" ")
        if k == "*RST":
            self.state = dict(self.factory)
            self.levels = {i: "0.000000" for i in self.levels}
            self.running = True
            return None
        if k == ":SYSTem:SETup":
            tok = v[3:] if v.startswith("#19") else None
            if tok in self.setups:
                st, lv, run = self.setups[tok]
                self.state, self.levels, self.running = dict(st), dict(lv), run
            else:
                self.errors.append('-222,"Data out of range"')
            return None
        if k in (":MEASure:DELete", ":MEASure:CLEar"):
            self.measurements_cleared = True
            return None
        if k == ":RUN":
            self.running = True
        elif k == ":STOP":
            self.running = False
        elif k == ":COUNter:TOTalize:CLEar":
            self.clear_t = self.clock.now()
        elif k == ":TRIGger:EDGE:LEVel":
            src = int(self.state[":TRIGger:EDGE:SOURce"][-1])
            if self.firmware == "mod" and not self.shown(src):
                return None  # the mod ignores a hidden channel's trigger level (seen on our bench, 2026-10-16)
            self.levels[int(self.state[":TRIGger:EDGE:SOURce"][-1])] = "%.6f" % float(v)
        elif k in self.state:
            m = re.match(r"^:CHANnel(\d):(SCALe|OFFSet)$", k)
            if m and self.firmware == "mod" and not self.shown(int(m.group(1))):
                return None  # the mod ignores a hidden channel's scale / offset (no error)
            self.state[k] = self.norm_value(k, v)
        else:
            self.errors.append('-113,"Undefined header"')
        return None

    def unknown(self):
        self.errors.append('-100,"Command err"')
        return None

    def snapshot(self):
        with self.lock:
            return dict(self.state), dict(self.levels), self.running


def read_json(path):
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def read_lines(path):
    with open(path, encoding="utf-8") as f:
        return f.read().splitlines()


def writes_of(lines):
    return [x for x in lines if "?" not in x and not x.startswith("*")]


class KitCase(unittest.TestCase):
    def setUp(self):
        self.clock = FakeClock()
        self.tmp = tempfile.mkdtemp(prefix="rate_kit_test_")
        self.scopes = []

    def tearDown(self):
        for s in self.scopes:
            s.close()
        shutil.rmtree(self.tmp, ignore_errors=True)

    def scope(self, **kw):
        s = FakeScope(self.clock, **kw)
        self.scopes.append(s)
        return s

    def measure(self, scope, *extra):
        argv = ["measure", "--host", scope.addr, "--timeout", "0.5", "--quiet", "--out", self.tmp,
                "--tdiv", "2e-8", "--depth", "1k", "--repeats", "1", "--seconds", "10"] + list(extra)
        code = rk.main(argv, clock=self.clock)
        js = [f for f in os.listdir(self.tmp) if f.startswith("rate-kit-") and f.endswith(".json")]
        rep = read_json(os.path.join(self.tmp, js[0])) if js else None
        return code, rep


class Identity(KitCase):
    def ident(self, **kw):
        s = self.scope(**kw)
        c = rk.Scpi("127.0.0.1", s.port, timeout=0.5, clock=self.clock)
        try:
            return rk.identify(c), s
        finally:
            c.close()

    def test_stock_mho(self):
        i, s = self.ident(firmware="stock")
        self.assertEqual((i.family, i.series, i.firmware, i.channels), ("mho", "MHO900", "rigol-stock", 4))
        self.assertIn(":NOVAos:VERSion?", s.received)
        self.assertEqual(s.errors, [], "the -100 of the unknown query is drained")

    def test_kiszka_mod_is_not_asked_novaos(self):
        i, s = self.ident(firmware="mod")
        self.assertEqual(i.firmware, "nk-mod")
        self.assertNotIn(":NOVAos:VERSion?", s.received)

    def test_novaos(self):
        i, _ = self.ident(firmware="novaos")
        self.assertEqual(i.firmware, "novaos")
        self.assertTrue(i.novaos_version.startswith("r81"))

    def test_dho_models(self):
        i, _ = self.ident(model="DHO924S")
        self.assertEqual((i.family, i.series, i.channels, i.default_generator()), ("dho", "DHO900", 4, "SOURce"))
        i, _ = self.ident(model="DHO802")
        self.assertEqual((i.series, i.channels), ("DHO800", 2))

    def test_other_models_refused(self):
        with self.assertRaises(rk.KitError):
            rk.Identity("RIGOL TECHNOLOGIES,DS1054Z,X,00.04", None)


class Pure(unittest.TestCase):
    def test_nrf_as_the_bench_engine(self):
        self.assertEqual([rk.nrf(x) for x in (1.5, 2e-8, -1.5, 0, 2e-6, 10e6)], ["1.5E0", "2E-8", "-1.5E0", "0", "2E-6", "1E7"])
        self.assertEqual([rk.trim_num(x) for x in (3.0, 1e7, 0.5)], ["3", "10000000", "0.5"])

    def test_depth_word(self):
        self.assertEqual([rk.depth_word(x) for x in ("1.0000E+03", "1.0000E+04", "1.0000E+06", "AUTO")], ["1k", "10k", "1M", "AUTO"])

    def test_settable_and_same(self):
        self.assertEqual(rk.settable(":CHANnel1:DISPlay", "1"), "ON")
        self.assertEqual(rk.settable(":COUNter:SOURce", "CHAN1"), "CHAN1")
        self.assertTrue(rk.same(":CHANnel1:SCALe", "5.000000E-01", "0.5"))
        self.assertTrue(rk.same(":COUNter:ENABle", "1", "ON"))

    def test_refuses_screenshot_and_reset(self):
        s = rk.Scpi("127.0.0.1", 9, timeout=0.1)
        for c in (":DISPlay:DATA? PNG", ":DISP:DATA?", "*RST", ":SAVE:IMAGe C:/x.png", ":SYSTem:RESet"):
            with self.assertRaises(rk.KitError):
                s.query(c)
        self.assertEqual(s.sent, [])


class Totalize(KitCase):
    def test_rate_is_count_difference_over_reply_times(self):
        s = self.scope(firmware="stock", rate=7345.0)
        code, rep = self.measure(s, "--method", "aux-self")
        self.assertEqual(code, 0, rep and rep.get("failed"))
        r = rep["scenes"][0]["repeats"][0]
        self.assertAlmostEqual(r["rate"], 7345.0, delta=1.0)
        self.assertAlmostEqual((r["count1"] - r["count0"]) / r["seconds"], r["rate"], delta=0.1)
        self.assertGreater(r["count0"], 0, "the first read comes --clear-wait after the clear")
        self.assertEqual(len(r["polls"]), 10)
        self.assertTrue(rep["restore"]["uut"]["complete"])

    def test_poll_counter_gives_a_slope(self):
        s = self.scope(rate=6000.0)
        code, rep = self.measure(s, "--method", "aux-self", "--poll", "counter")
        r = rep["scenes"][0]["repeats"][0]
        self.assertAlmostEqual(r["rate_fit"], 6000.0, delta=2.0)
        self.assertGreater(r["distinct_counts"], 2)

    def test_repeats_summary(self):
        s = self.scope(rate=5000.0)
        code, rep = self.measure(s, "--method", "aux-self", "--repeats", "3")
        sm = rep["scenes"][0]["summary"]["aux_rate"]
        self.assertEqual(sm["n"], 3)
        self.assertAlmostEqual(sm["mean"], 5000.0, delta=1.0)

    def test_srate_reported_with_the_counter_channel(self):
        s = self.scope(rate=5000.0)
        code, rep = self.measure(s, "--method", "aux-self")
        self.assertEqual(rep["scenes"][0]["srate"], "2.000000E+09")


class HiddenChannel(KitCase):
    def test_mod_hidden_counts_when_set_up_while_shown(self):
        # CH4 hidden at the start: the mod would ignore its scale, offset and level, so the kit shows it first
        s = self.scope(firmware="mod", rate=9200.0, seed={":CHANnel4:DISPlay": "0"})
        before = s.snapshot()
        code, rep = self.measure(s, "--method", "aux-self", "--counter-channel", "hidden")
        sc = rep["scenes"][0]
        self.assertAlmostEqual(sc["repeats"][0]["rate"], 9200.0, delta=1.0)
        self.assertNotIn("warning", sc)
        self.assertTrue(rep["restore"]["uut"]["complete"], rep["restore"])
        self.assertEqual(s.snapshot()[1], before[1], "the levels come back although CH4 ends hidden")

    def test_zero_threshold_counts_nothing_and_says_why(self):
        s = self.scope(firmware="stock")
        s.aux = 3  # the pulses arrive on CH3 while the counter counts CH4
        code, rep = self.measure(s, "--method", "aux-self")
        self.assertIn("threshold", rep["scenes"][0]["warning"])

    def test_mod_shown_counts(self):
        s = self.scope(firmware="mod", rate=6761.0)
        code, rep = self.measure(s, "--method", "aux-self")
        self.assertEqual(code, 0)
        self.assertAlmostEqual(rep["scenes"][0]["repeats"][0]["rate"], 6761.0, delta=1.0)
        self.assertIn(":CHANnel4:DISPlay ON", s.received)

    def test_stock_counts_hidden(self):
        s = self.scope(firmware="stock", rate=12033.0)
        code, rep = self.measure(s, "--method", "aux-self", "--counter-channel", "hidden")
        self.assertAlmostEqual(rep["scenes"][0]["repeats"][0]["rate"], 12033.0, delta=1.0)
        self.assertNotIn("warning", rep["scenes"][0])

    def test_hidden_aux_input_is_shown_before_its_scale(self):
        s = self.scope(firmware="stock")
        self.measure(s, "--method", "aux-self", "--counter-channel", "hidden")
        w = writes_of(s.received)
        self.assertLess(w.index(":CHANnel4:DISPlay ON"), w.index(":CHANnel4:SCALe 1"))


class Restore(KitCase):
    def check_restored(self, s, before):
        st, lv, run = s.snapshot()
        b_st, b_lv, b_run = before
        diff = {k: (b_st[k], st[k]) for k in b_st if not rk.same(k, st.get(k), b_st[k])}
        self.assertEqual(diff, {})
        self.assertEqual(lv, b_lv)
        self.assertEqual(run, b_run)

    def test_restore_on_success(self):
        s = self.scope(seed={":CHANnel2:DISPlay": "1", ":CHANnel2:SCALe": "5.000000E-01"})
        before = s.snapshot()
        code, rep = self.measure(s, "--method", "aux-self", "--channels", "one,all", "--depth", "1k,AUTO")
        self.assertEqual(code, 0, rep["failed"])
        self.check_restored(s, before)

    def test_restore_on_error(self):
        s = self.scope()
        s.never_answer.add(":COUNter:CURRent")
        before = s.snapshot()
        code, rep = self.measure(s, "--method", "aux-self")
        self.assertEqual(code, 1)
        self.assertIn("COUNter:CURRent", rep["failed"])
        self.assertTrue(rep["restore"]["uut"]["complete"])
        self.check_restored(s, before)

    def test_restore_on_ctrl_c(self):
        s = self.scope(seed={":CHANnel4:DISPlay": "1"})
        s.running = True
        before = s.snapshot()
        n = {"k": 0}

        def boom(sec):
            n["k"] += 1
            if n["k"] == 30:
                raise KeyboardInterrupt

        self.clock.on_sleep = boom
        code, rep = self.measure(s, "--method", "aux-self")
        self.clock.on_sleep = None
        self.assertEqual(code, 130)
        self.check_restored(s, before)

    def test_hung_scpi_is_reported_and_the_snapshot_kept(self):
        s = self.scope()
        s.hang_after = ":COUNter:TOTalize:CLEar"
        code, rep = self.measure(s, "--method", "aux-self")
        self.assertEqual(code, 4)
        self.assertFalse(rep["restore"]["uut"]["complete"])
        self.assertIn("rate_kit.py restore", rep["restore"]["uut"]["why"])
        snaps = [f for f in os.listdir(self.tmp) if f.startswith("snapshot-")]
        self.assertEqual(len(snaps), 1)
        # once it answers again, restore from the snapshot puts it back
        before_vals = read_json(os.path.join(self.tmp, snaps[0]))["uut"]["snapshot"]["values"]
        s.hung = False
        s.hang_after = None
        code = rk.main(["restore", "--host", s.addr, "--timeout", "0.5", "--from", os.path.join(self.tmp, snaps[0])],
                       clock=self.clock)
        self.assertEqual(code, 0)
        st, _, _ = s.snapshot()
        for k, v in before_vals.items():
            if k == ":TRIGger:EDGE:LEVel":
                continue  # kept per source in the fake
            self.assertTrue(rk.same(k, st[k], v), k)

    def test_no_measurement_unless_asked(self):
        s = self.scope()
        self.measure(s, "--method", "aux-self")
        self.assertFalse([x for x in s.received if x.startswith(":MEASure")])

    def test_never_sends_screenshot(self):
        s = self.scope(firmware="mod")
        self.measure(s, "--method", "aux-self", "--display", "bench")
        self.assertFalse([x for x in s.received if "DATA" in x.upper()])


class Redaction(KitCase):
    def test_serial_and_address_hidden(self):
        s = self.scope(serial="MHO9X000T00002")
        self.tmp = os.path.join(self.tmp, "MHO9X000T00002")  # an --out named after the serial lands in argv
        os.makedirs(self.tmp)
        code, rep = self.measure(s, "--method", "aux-self")
        blob = "".join(" ".join(read_lines(os.path.join(self.tmp, f))) for f in os.listdir(self.tmp))
        self.assertNotIn("MHO9X000T00002", blob)
        self.assertNotIn("127.0.0.1", blob)
        self.assertIn("<serial hidden>", rep["scopes"]["uut"]["idn"])

    def test_show_serial(self):
        s = self.scope(serial="MHO9X000T00002")
        code, rep = self.measure(s, "--method", "aux-self", "--show-serial")
        self.assertEqual(rep["scopes"]["uut"]["serial"], "MHO9X000T00002")


class External(KitCase):
    def test_counter_on_another_scope(self):
        uut = self.scope(firmware="stock", rate=0.0)
        cnt = self.scope(firmware="stock", model="DHO924S", rate=11000.0, aux=1)
        before_u, before_c = uut.snapshot(), cnt.snapshot()
        code, rep = self.measure(uut, "--method", "aux-external", "--counter", cnt.addr)
        self.assertEqual(code, 0, rep["failed"])
        self.assertAlmostEqual(rep["scenes"][0]["repeats"][0]["rate"], 11000.0, delta=1.0)
        self.assertFalse([x for x in writes_of(uut.received) if x.startswith(":COUNter")], "the UUT's counter is not used")
        self.assertIn(":SYSTem:AOUTput TOUT", uut.received)
        self.assertEqual(rep["scenes"][0]["srate"], "4.000000E+09", "the UUT runs its signal channel alone")
        for s, b in ((uut, before_u), (cnt, before_c)):
            st, lv, run = s.snapshot()
            self.assertEqual({k: (b[0][k], st[k]) for k in b[0] if not rk.same(k, st[k], b[0][k])}, {})


class MoreOptions(KitCase):
    def test_poll_quiet_sends_nothing_in_the_window(self):
        s = self.scope(rate=8000.0)
        code, rep = self.measure(s, "--method", "aux-self", "--poll", "quiet")
        r = rep["scenes"][0]["repeats"][0]
        self.assertEqual(r["polls"], [])
        self.assertAlmostEqual(r["rate"], 8000.0, delta=1.0)
        i0 = max(i for i, x in enumerate(s.received) if x == ":COUNter:TOTalize:CLEar")
        between = s.received[i0 + 1:]
        first = between.index(":COUNter:CURRent?")
        second = between.index(":COUNter:CURRent?", first + 1)
        self.assertEqual(between[first + 1:second], [])

    def test_display_default_and_function(self):
        s = self.scope()
        self.measure(s, "--method", "aux-self", "--display", "default", "--function", "SINusoid")
        w = writes_of(s.received)
        self.assertIn(":DISPlay:GRADing:TIME MIN", w)
        self.assertIn(":DISPlay:COLor OFF", w)
        self.assertIn(":SOURce2:FUNCtion SINusoid", w)

    def test_clear_measurements_only_when_asked(self):
        s = self.scope()
        code, rep = self.measure(s, "--method", "aux-self", "--clear-measurements")
        self.assertIn(":MEASure:DELete", s.received)
        self.assertTrue(any("cleared" in n for n in rep["notes"]))


class Health(KitCase):
    def run_health(self, s, *extra):
        code = rk.main(["scpi-health", "--host", s.addr, "--timeout", "0.5", "--quiet", "--out", self.tmp, "--n", "3"]
                       + list(extra), clock=self.clock)
        js = [f for f in os.listdir(self.tmp) if f.startswith("scpi-health-") and f.endswith(".json")]
        return code, read_json(os.path.join(self.tmp, js[0]))

    def test_healthy(self):
        s = self.scope()
        code, rep = self.run_health(s)
        self.assertEqual(code, 0)
        self.assertEqual(rep["unanswered"], {"*IDN?": 0, ":TRIGger:STATus?": 0})
        self.assertFalse([x for x in s.received if "DATA" in x.upper()])


class Preset(KitCase):
    def test_reset_count_and_setup_back(self):
        s = self.scope(firmware="stock", rate=25000.0, seed={":CHANnel2:DISPlay": "1", ":CHANnel1:SCALe": "2.000000E-01"})
        s.lenient = True
        before = s.snapshot()
        code = rk.main(["preset", "--host", s.addr, "--timeout", "0.5", "--quiet", "--out", self.tmp, "--method",
                        "aux-self", "--tdiv", "2e-8", "--depth", "1k", "--repeats", "1", "--yes-factory-default"],
                       clock=self.clock)
        self.assertEqual(code, 0)
        js = [f for f in os.listdir(self.tmp) if f.startswith("preset-") and f.endswith(".json")]
        rep = read_json(os.path.join(self.tmp, js[0]))
        self.assertIn("*RST", s.received)
        self.assertIn(":CHANnel1:SCALe", rep["reset_changed"])
        self.assertEqual(rep["restore_diff"], {})
        self.assertAlmostEqual(rep["measure_summary"][0]["summary"]["aux_rate"]["mean"], 25000.0, delta=2.0)
        st, lv, run = s.snapshot()
        self.assertEqual({k: v for k, v in st.items() if not rk.same(k, v, before[0][k])}, {})
        self.assertEqual(run, before[2])
        # --scene minimal is the default here: no probe / coupling writes
        self.assertFalse([x for x in s.received if x.startswith(":CHANnel1:PROBe ")])

    def test_dho_is_not_asked_the_input_impedance(self):
        # the DHO800/DHO900 guide has no :CHANnel<n>:IMPedance (1 MOhm inputs only); the MHO900 guide has it (3.6.7)
        s = self.scope(firmware="stock", model="DHO924S", rate=20000.0, aux=3)
        s.lenient = True
        code = rk.main(["preset", "--host", s.addr, "--timeout", "0.5", "--quiet", "--out", self.tmp, "--method",
                        "aux-self", "--tdiv", "2e-8", "--depth", "1k", "--repeats", "1", "--yes-factory-default"],
                       clock=self.clock)
        self.assertEqual(code, 0)
        self.assertFalse([x for x in s.received if "IMPedance" in x])

    def test_needs_the_flag(self):
        s = self.scope()
        code = rk.main(["preset", "--host", s.addr, "--quiet", "--out", self.tmp], clock=self.clock)
        self.assertEqual(code, 2)
        self.assertNotIn("*RST", s.received)


class Trio(KitCase):
    def test_one_table_for_three_units(self):
        st = self.scope(firmware="stock", rate=25000.0)
        mo = self.scope(firmware="mod", rate=15000.0)
        nv = self.scope(firmware="novaos", rate=46000.0)
        code = rk.main(["trio", "--unit", "stock=" + st.addr, "--unit", "mod=" + mo.addr, "--unit", "novaos=" + nv.addr,
                        "--timeout", "0.5", "--quiet", "--out", self.tmp, "--method", "aux-self", "--tdiv", "2e-8",
                        "--depth", "1k", "--repeats", "1"], clock=self.clock)
        self.assertEqual(code, 0)
        md = [f for f in os.listdir(self.tmp) if f.startswith("trio-") and f.endswith(".md")]
        text = " ".join(read_lines(os.path.join(self.tmp, md[0])))
        for lab in ("stock", "mod", "novaos"):
            self.assertTrue(os.path.isdir(os.path.join(self.tmp, lab)))
            self.assertIn(lab, text)
        self.assertIn("AUX 25,000", text)
        self.assertIn("AUX 15,000", text)
        self.assertIn("engine 46,000", text)
        self.assertNotIn("127.0.0.1", text)

    def test_load_counter_on_the_unit(self):
        uut = self.scope(firmware="stock", rate=0.0)
        cnt = self.scope(firmware="stock", model="DHO924S", rate=11000.0, aux=1)
        before = uut.snapshot()
        code, rep = self.measure(uut, "--method", "aux-external", "--counter", cnt.addr, "--load-counter-input", "4")
        self.assertEqual(code, 0, rep["failed"])
        w = writes_of(uut.received)
        self.assertIn(":COUNter:SOURce CHANnel4", w)
        self.assertLess(w.index(":CHANnel4:SCALe 1"), w.index(":CHANnel4:DISPlay OFF"))
        self.assertEqual(rep["scenes"][0]["srate"], "2.000000E+09", "the hidden counter channel is digitised")
        st, _, _ = uut.snapshot()
        self.assertEqual({k: v for k, v in st.items() if not rk.same(k, v, before[0][k])}, {})

    def test_unit_spec(self):
        u = rk.parse_unit("stock=stock.lan,method=aux-external,counter=counter.lan@3")
        self.assertEqual((u["label"], u["host"], u["method"], u["counter"]), ("stock", "stock.lan", "aux-external", "counter.lan@3"))


if __name__ == "__main__":
    unittest.main(verbosity=2)