use std::collections::BTreeMap;
use std::time::Duration;
use crate::api::explore::v1::{GoalCandidates, State as ExploreState};
use crate::api::follow::v1::State as FollowState;
use crate::api::localize::v1::LocalizationState;
use crate::api::map::v1::{Summary as MapSummary, TraversabilitySummary};
use crate::api::mission::v1::{
ExplorationCompletion, ExplorationCompletionMode, GoalPose, GoalTolerance, MissionCommand,
State as MissionState,
};
use crate::api::plan::v1::State as PlanState;
use crate::api::presence::Summary;
use crate::api::safety::v1::State as SafetyState;
use crate::api::simulation::v1::{
clock, clock::Clock, pose, pose::Pose, reset, status, status::Status,
};
use crate::bus::Bus;
use crate::bus::pubsub::Stamped;
use crate::bus::zenoh::{TypedPublisher, TypedSchema, TypedSubscriber};
use crate::runtime::DEFAULT_ROBOT_NAMESPACE;
use anyhow::{Context, Result, anyhow, bail};
use serde::{Serialize, de::DeserializeOwned};
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct ScenarioEnvironment {
pub robot_router_endpoint: String,
pub robot_namespace: String,
pub robot_id: String,
}
pub struct ScenarioContext {
bus: Bus,
environment: ScenarioEnvironment,
wallclock_timeout: Duration,
}
impl ScenarioEnvironment {
pub fn from_env() -> Result<Self> {
Self::from_vars(|key| std::env::var(key))
}
pub fn from_map(vars: &BTreeMap<String, String>) -> Result<Self> {
Self::from_vars(|key| vars.get(key).cloned().ok_or(std::env::VarError::NotPresent))
}
fn from_vars(
var: impl Fn(&str) -> std::result::Result<String, std::env::VarError>,
) -> Result<Self> {
let robot_router_endpoint =
var(crate::runtime::ENV_ROBOT_ROUTER_ENDPOINT).with_context(|| {
format!(
"{} must be set for scenario context",
crate::runtime::ENV_ROBOT_ROUTER_ENDPOINT
)
})?;
let robot_id = var(crate::runtime::ENV_ROBOT_ID).with_context(|| {
format!(
"{} must be set for scenario context",
crate::runtime::ENV_ROBOT_ID
)
})?;
let robot_namespace = var(crate::runtime::ENV_ROBOT_NAMESPACE)
.unwrap_or_else(|_| DEFAULT_ROBOT_NAMESPACE.to_string());
Self::new(robot_router_endpoint, robot_namespace, robot_id)
}
pub fn new(
robot_router_endpoint: impl Into<String>,
robot_namespace: impl Into<String>,
robot_id: impl Into<String>,
) -> Result<Self> {
let environment = Self {
robot_router_endpoint: trim_required(
robot_router_endpoint.into(),
crate::runtime::ENV_ROBOT_ROUTER_ENDPOINT,
)?,
robot_namespace: trim_required(
robot_namespace.into(),
crate::runtime::ENV_ROBOT_NAMESPACE,
)?,
robot_id: trim_required(robot_id.into(), crate::runtime::ENV_ROBOT_ID)?,
};
Ok(environment)
}
pub async fn connect(self, wallclock_timeout: Duration) -> Result<ScenarioContext> {
let bus = crate::bus::builder::Builder::new(self.robot_router_endpoint.clone())
.with_prefix(self.robot_namespace.clone())
.connect()
.await?;
Ok(ScenarioContext {
bus,
environment: self,
wallclock_timeout,
})
}
}
impl ScenarioContext {
pub async fn from_env() -> Result<Self> {
ScenarioEnvironment::from_env()
.context("failed to build scenario environment from process env")?
.connect(Duration::from_secs(30))
.await
}
pub fn bus(&self) -> &Bus {
&self.bus
}
pub fn environment(&self) -> &ScenarioEnvironment {
&self.environment
}
pub async fn reset_simulation(&self) -> Result<u64> {
let retry =
crate::bus::query::Retry::new(3).with_initial_backoff(Duration::from_millis(50));
let response = reset::query(&self.bus, &reset::Request, &retry)
.await?
.ok_or_else(|| anyhow!("simulation reset returned no acknowledgement"))?;
self.wait_for_status(|status| status.epoch >= response.epoch && status.step > 0)
.await?;
Ok(response.epoch)
}
pub async fn wait_until_ready(&self) -> Result<Stamped<Status>> {
self.wait_for_status(|status| status.step > 0).await
}
pub async fn advance_for_secs(&self, secs: f64) -> Result<Stamped<Clock>> {
let duration_ns = duration_ns_from_secs(secs)?;
let subscriber = clock::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
let first = next_stamped(&subscriber, self.wallclock_timeout).await?;
let epoch = first.data.epoch();
let target_time_ns = first
.data
.time_ns()
.checked_add(duration_ns)
.ok_or_else(|| anyhow!("scenario target logical time overflows nanoseconds"))?;
let mut latest = first;
while latest.data.time_ns() < target_time_ns {
latest = next_stamped(&subscriber, self.wallclock_timeout).await?;
if latest.data.epoch() != epoch {
bail!(
"simulation epoch changed from {} to {} while waiting for logical time",
epoch,
latest.data.epoch()
);
}
}
Ok(latest)
}
pub async fn simulation_pose(&self) -> Result<Stamped<Pose>> {
let subscriber = pose::subscriber_builder(&self.bus, &self.environment.robot_id)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_localization_state(&self) -> Result<Stamped<LocalizationState>> {
let subscriber = crate::api::localize::v1::state::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_presence_summary(&self) -> Result<Stamped<Summary>> {
let subscriber = crate::api::presence::summary::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_safety_state(&self) -> Result<Stamped<SafetyState>> {
let subscriber = crate::api::safety::v1::state::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_plan_state(&self) -> Result<Stamped<PlanState>> {
let subscriber = crate::api::plan::v1::state::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_follow_state(&self) -> Result<Stamped<FollowState>> {
let subscriber = crate::api::follow::v1::state::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_mission_state(&self) -> Result<Stamped<MissionState>> {
let subscriber = crate::api::mission::v1::state::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_explore_state(&self) -> Result<Stamped<ExploreState>> {
let subscriber = crate::api::explore::v1::state::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_explore_candidates(&self) -> Result<Stamped<GoalCandidates>> {
let subscriber = crate::api::explore::v1::goal_candidates::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_map_summary(&self) -> Result<Stamped<MapSummary>> {
let subscriber = crate::api::map::v1::summary::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn latest_traversability_summary(&self) -> Result<Stamped<TraversabilitySummary>> {
let subscriber = crate::api::map::v1::traversability_summary::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
next_stamped(&subscriber, self.wallclock_timeout).await
}
pub async fn publish_navigate_to(
&self,
goal: GoalPose,
tolerance: GoalTolerance,
) -> Result<()> {
self.publish_mission_command(MissionCommand::NavigateTo {
goal,
tolerance,
max_duration_ns: None,
})
.await
}
pub async fn publish_explore_command(&self) -> Result<()> {
self.publish_mission_command(MissionCommand::Explore {
area: None,
completion: ExplorationCompletion {
mode: ExplorationCompletionMode::OpenEnded,
coverage_goal: None,
},
max_duration_ns: None,
})
.await
}
pub async fn publish_cancel(&self) -> Result<()> {
self.publish_mission_command(MissionCommand::Cancel).await
}
pub async fn publish_pause(&self) -> Result<()> {
self.publish_mission_command(MissionCommand::Pause).await
}
pub async fn publish_manual_command(
&self,
command: crate::api::motion::v1::ManualCommand,
) -> Result<()> {
let produced_at_ns = self.wait_until_ready().await?.data.time_ns;
let publisher = crate::api::motion::v1::manual::publisher(&self.bus)?
.await
.map_err(|error| anyhow!(error.to_string()))?;
wait_for_matching_subscriber(&publisher, self.wallclock_timeout).await?;
publisher
.put(&Stamped::new(produced_at_ns, command))
.await
.map_err(|error| anyhow!(error.to_string()))
}
async fn wait_for_status(
&self,
predicate: impl Fn(&Status) -> bool,
) -> Result<Stamped<Status>> {
let subscriber = status::subscriber_builder(&self.bus)
.await
.map_err(|error| anyhow!(error.to_string()))?;
loop {
let status = next_stamped(&subscriber, self.wallclock_timeout).await?;
if predicate(&status.data) {
return Ok(status);
}
}
}
async fn publish_mission_command(&self, command: MissionCommand) -> Result<()> {
let produced_at_ns = self.wait_until_ready().await?.data.time_ns;
let publisher = crate::api::mission::v1::command::publisher(&self.bus)?
.await
.map_err(|error| anyhow!(error.to_string()))?;
wait_for_matching_subscriber(&publisher, self.wallclock_timeout).await?;
publisher
.put(&Stamped::new(produced_at_ns, command))
.await
.map_err(|error| anyhow!(error.to_string()))
}
}
async fn next_stamped<T>(
subscriber: &TypedSubscriber<Stamped<T>>,
timeout: Duration,
) -> Result<Stamped<T>>
where
T: DeserializeOwned + TypedSchema,
{
match tokio::time::timeout(timeout, subscriber.recv_async()).await {
Ok(Ok(Ok(value))) => Ok(value),
Ok(Ok(Err(error))) => Err(anyhow!("typed scenario payload failed to decode: {error}")),
Ok(Err(error)) => Err(anyhow!("scenario subscriber failed: {error}")),
Err(_) => bail!("timed out waiting for scenario data after {:?}", timeout),
}
}
async fn wait_for_matching_subscriber<T>(
publisher: &TypedPublisher<'_, T>,
timeout: Duration,
) -> Result<()>
where
T: Serialize + TypedSchema,
{
let deadline = tokio::time::Instant::now() + timeout;
loop {
if publisher
.has_matching_subscribers()
.await
.map_err(|error| anyhow!(error.to_string()))?
{
return Ok(());
}
if tokio::time::Instant::now() >= deadline {
bail!("no subscriber matched the scenario command publisher within {timeout:?}");
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
fn trim_required(value: String, name: &str) -> Result<String> {
let value = value.trim().to_string();
if value.is_empty() {
bail!("{name} must not be empty");
}
Ok(value)
}
fn duration_ns_from_secs(secs: f64) -> Result<u64> {
let duration = Duration::try_from_secs_f64(secs).map_err(|_| {
anyhow!("scenario advance duration must be finite and non-negative, got {secs}")
})?;
duration
.as_nanos()
.try_into()
.map_err(|_| anyhow!("scenario advance duration overflows nanoseconds: {secs} seconds"))
}
#[cfg(test)]
mod tests {
use super::{ScenarioEnvironment, duration_ns_from_secs};
use std::collections::BTreeMap;
#[test]
fn environment_uses_existing_runtime_env_names() -> anyhow::Result<()> {
let vars = BTreeMap::from([
(
crate::runtime::ENV_ROBOT_ROUTER_ENDPOINT.to_string(),
"tcp/127.0.0.1:7447".to_string(),
),
(
crate::runtime::ENV_ROBOT_NAMESPACE.to_string(),
"dev".to_string(),
),
(
crate::runtime::ENV_ROBOT_ID.to_string(),
"robot-a".to_string(),
),
]);
let env = ScenarioEnvironment::from_map(&vars)?;
assert_eq!(env.robot_router_endpoint, "tcp/127.0.0.1:7447");
assert_eq!(env.robot_namespace, "dev");
assert_eq!(env.robot_id, "robot-a");
Ok(())
}
#[test]
fn environment_requires_router_endpoint() {
let vars = BTreeMap::from([(
crate::runtime::ENV_ROBOT_ID.to_string(),
"robot-a".to_string(),
)]);
assert!(ScenarioEnvironment::from_map(&vars).is_err());
}
#[test]
fn duration_accepts_fractional_seconds() -> anyhow::Result<()> {
assert_eq!(duration_ns_from_secs(1.5)?, 1_500_000_000);
Ok(())
}
#[test]
fn duration_rejects_invalid_seconds() {
assert!(duration_ns_from_secs(f64::NAN).is_err());
assert!(duration_ns_from_secs(-1.0).is_err());
}
}