ria-toolkit-oss/tests/orchestration/test_labeler.py

146 lines
5.6 KiB
Python

"""Tests for orchestration labeler."""
import time
import numpy as np
import pytest
from ria_toolkit_oss.datatypes.recording import Recording
from ria_toolkit_oss.orchestration.campaign import CaptureStep
from ria_toolkit_oss.orchestration.labeler import build_output_filename, label_recording
def _simple_recording() -> Recording:
sr = 1e6
n = 1000
data = np.ones(n, dtype=np.complex64)
return Recording(data, metadata={"sample_rate": sr, "center_frequency": 2.45e9})
def _wifi_step() -> CaptureStep:
return CaptureStep(
duration=30.0,
label="ch06_20mhz_idle",
channel=6,
bandwidth_mhz=20.0,
traffic="idle",
)
def _bt_step() -> CaptureStep:
return CaptureStep(
duration=30.0,
label="audio_stream",
traffic="audio_stream",
connection_interval_ms=7.5,
)
# ---------------------------------------------------------------------------
# label_recording
# ---------------------------------------------------------------------------
class TestLabelRecording:
def test_device_id_set(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert rec.metadata["device_id"] == "iphone13_001"
def test_capture_timestamp_set(self):
ts = time.time()
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), ts)
assert rec.metadata["capture_timestamp"] == pytest.approx(ts, abs=1.0)
def test_step_label_set(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert rec.metadata["step_label"] == "ch06_20mhz_idle"
def test_step_duration_set(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert rec.metadata["step_duration_s"] == pytest.approx(30.0)
def test_campaign_name_optional(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert "campaign" not in rec.metadata
def test_campaign_name_when_provided(self):
rec = label_recording(
_simple_recording(), "iphone13_001", _wifi_step(), time.time(), campaign_name="test_campaign"
)
assert rec.metadata["campaign"] == "test_campaign"
def test_wifi_channel_set(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert rec.metadata["wifi_channel"] == 6
def test_wifi_bandwidth_set(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert rec.metadata["wifi_bandwidth_mhz"] == pytest.approx(20.0)
def test_traffic_pattern_set(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert rec.metadata["traffic_pattern"] == "idle"
def test_bt_connection_interval_set(self):
rec = label_recording(_simple_recording(), "airpods_001", _bt_step(), time.time())
assert rec.metadata["bt_connection_interval_ms"] == pytest.approx(7.5)
def test_no_channel_key_for_bt(self):
"""BT steps with no channel should not add wifi_channel to metadata."""
rec = label_recording(_simple_recording(), "airpods_001", _bt_step(), time.time())
assert "wifi_channel" not in rec.metadata
def test_no_bandwidth_key_for_bt(self):
rec = label_recording(_simple_recording(), "airpods_001", _bt_step(), time.time())
assert "wifi_bandwidth_mhz" not in rec.metadata
def test_power_dbm_set(self):
step = CaptureStep(duration=30.0, label="test", traffic="idle", power_dbm=15.0)
rec = label_recording(_simple_recording(), "dev_001", step, time.time())
assert rec.metadata["tx_power_dbm"] == pytest.approx(15.0)
def test_no_power_key_when_unset(self):
rec = label_recording(_simple_recording(), "iphone13_001", _wifi_step(), time.time())
assert "tx_power_dbm" not in rec.metadata
def test_returns_same_recording(self):
"""label_recording should mutate and return the same Recording object."""
rec = _simple_recording()
result = label_recording(rec, "iphone13_001", _wifi_step(), time.time())
assert result is rec
# ---------------------------------------------------------------------------
# build_output_filename
# ---------------------------------------------------------------------------
class TestBuildOutputFilename:
def test_basic_wifi(self):
step = CaptureStep(duration=30.0, label="ch06_20mhz_idle")
fn = build_output_filename("iphone13_wifi_001", step)
assert fn == "iphone13_wifi_001/ch06_20mhz_idle"
def test_bt_step(self):
step = CaptureStep(duration=30.0, label="audio_stream")
fn = build_output_filename("airpods_pro_bt_001", step)
assert fn == "airpods_pro_bt_001/audio_stream"
def test_spaces_in_device_id_replaced(self):
step = CaptureStep(duration=30.0, label="idle")
fn = build_output_filename("my device", step)
assert " " not in fn
assert fn == "my_device/idle"
def test_slashes_in_label_replaced(self):
step = CaptureStep(duration=30.0, label="ch/6/idle")
fn = build_output_filename("dev_001", step)
assert "/" not in fn.split("/", 1)[1] # only the separator slash should remain
def test_path_structure(self):
"""Filename should be exactly '<device_id>/<label>' (one level of nesting)."""
step = CaptureStep(duration=30.0, label="idle")
fn = build_output_filename("dev_001", step)
parts = fn.split("/")
assert len(parts) == 2