use core::future::Future;
use std::pin::Pin;
use serenity::all::{
AutocompleteChoice, CommandInteraction, CommandOptionType, Context, CreateAutocompleteResponse,
CreateCommand, CreateCommandOption, CreateInteractionResponse,
CreateInteractionResponseFollowup, Permissions, ResolvedOption, ResolvedValue,
};
use shep_client::shep_core::protocol::{ProcessInfo, SelectorSpec};
use crate::{
bot::{
command::{Command, State},
commands::bare_subcommand,
embed,
},
error::Error,
limits,
shepherd::{Live, Verb},
stream::pack::pack_by_budget,
};
const AUTOCOMPLETE_CHOICE_LIMIT: usize = 25;
pub struct ShepCommand;
const VERBS: &[(&str, Verb)] = &[
("start", Verb::Start),
("stop", Verb::Stop),
("restart", Verb::Restart),
("reload", Verb::Reload),
("delete", Verb::Delete),
("flush", Verb::Flush),
("save", Verb::Save),
("reopen", Verb::Reopen),
];
fn verb_for_subcommand(name: &str) -> Option<Verb> {
VERBS
.iter()
.find(|(candidate, _)| *candidate == name)
.map(|(_, verb)| *verb)
}
fn name_option() -> CreateCommandOption {
CreateCommandOption::new(CommandOptionType::String, "name", "Which sheep, by name.")
.required(true)
.set_autocomplete(true)
}
fn optional_name_option() -> CreateCommandOption {
CreateCommandOption::new(CommandOptionType::String, "name", "Which sheep, by name.")
.required(false)
.set_autocomplete(true)
}
fn verb_subcommand(name: &str, description: &str) -> CreateCommandOption {
CreateCommandOption::new(CommandOptionType::SubCommand, name, description)
.add_sub_option(name_option())
}
fn autocomplete_failed_message(err: &serenity::Error) -> String {
format!("shep-discord: /shep autocomplete failed: {err}")
}
fn empty_flock_message() -> &'static str {
"There is nothing in the flock to list."
}
fn reply_content(reply: &str) -> String {
limits::fit(reply, limits::MESSAGE_CONTENT_LIMIT)
}
fn name_from(value: &ResolvedValue<'_>) -> Option<String> {
let ResolvedValue::SubCommand(options) = value else {
return None;
};
options.iter().find_map(|option| {
if option.name != "name" {
return None;
}
match option.value {
ResolvedValue::String(name) => Some(name.to_owned()),
_ => None,
}
})
}
fn selector_for_name(verb: Verb, name: Option<String>) -> SelectorSpec {
match verb {
Verb::Save => SelectorSpec::All,
Verb::Reopen => name.map_or(SelectorSpec::All, SelectorSpec::Name),
_ => SelectorSpec::Name(name.unwrap_or_default()),
}
}
fn selector_for(verb: Verb, value: &ResolvedValue<'_>) -> SelectorSpec {
selector_for_name(verb, name_from(value))
}
impl ShepCommand {
async fn suggestions(&self, live: &Live, partial: &str) -> Vec<String> {
let Ok(flock) = live.flock().await else {
return Vec::new();
};
let needle = partial.to_lowercase();
flock
.into_iter()
.map(|info| info.name)
.filter(|name| name.to_lowercase().contains(&needle))
.filter(|name| name.chars().count() <= limits::AUTOCOMPLETE_CHOICE_NAME_LIMIT)
.take(AUTOCOMPLETE_CHOICE_LIMIT)
.collect()
}
async fn flock_for_listing(
&self,
live: &Live,
ignore_dogs: bool,
) -> Result<Vec<ProcessInfo>, Error> {
let flock = live.flock().await?;
Ok(if ignore_dogs {
flock
.into_iter()
.filter(|info| info.dog.is_none())
.collect()
} else {
flock
})
}
async fn answer(
&self,
ctx: &Context,
interaction: &CommandInteraction,
content: &str,
) -> Result<(), Error> {
interaction
.create_followup(
ctx,
CreateInteractionResponseFollowup::new()
.content(content)
.ephemeral(true),
)
.await?;
Ok(())
}
async fn answer_list(
&self,
ctx: &Context,
interaction: &CommandInteraction,
state: &State,
) -> Result<(), Error> {
let flock = self.flock_for_listing(&state.live, true).await?;
if flock.is_empty() {
return self.answer(ctx, interaction, empty_flock_message()).await;
}
for group in pack_by_budget(
flock,
limits::MESSAGE_CHARACTER_BUDGET,
limits::EMBED_MAX_COUNT,
embed::embed_character_count,
) {
let embeds = group.iter().map(embed::process_embed).collect();
interaction
.create_followup(
ctx,
CreateInteractionResponseFollowup::new()
.embeds(embeds)
.ephemeral(true),
)
.await?;
}
Ok(())
}
async fn answer_verb(
&self,
verb: Verb,
selector: SelectorSpec,
ctx: &Context,
interaction: &CommandInteraction,
state: &State,
) -> Result<(), Error> {
let reply = state.live.act(verb, selector).await?;
self.answer(ctx, interaction, &reply_content(&reply)).await
}
}
impl Command for ShepCommand {
fn data(&self) -> CreateCommand {
CreateCommand::new(self.name())
.description("Drive the flock shep tends.")
.default_member_permissions(Permissions::ADMINISTRATOR)
.add_option(verb_subcommand(
"start",
"Restart a sheep that is already registered.",
))
.add_option(verb_subcommand("stop", "Stop a sheep, staying registered."))
.add_option(verb_subcommand("restart", "Restart a sheep."))
.add_option(verb_subcommand(
"reload",
"Replace a sheep with a fresh instance, one at a time.",
))
.add_option(verb_subcommand("delete", "Stop and deregister a sheep."))
.add_option(verb_subcommand("flush", "Empty a sheep's log files."))
.add_option(bare_subcommand("list", "List every sheep in the flock."))
.add_option(bare_subcommand("save", "Write the muster roll now."))
.add_option(
CreateCommandOption::new(
CommandOptionType::SubCommand,
"reopen",
"Reopen a sheep's log files, or every sheep's, for an external rotator.",
)
.add_sub_option(optional_name_option()),
)
}
fn name(&self) -> &'static str {
"shep"
}
fn run<'a>(
&'a self,
ctx: &'a Context,
interaction: &'a CommandInteraction,
state: &'a State,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
Box::pin(async move {
let options = interaction.data.options();
let sub: &ResolvedOption<'_> = match options.first() {
Some(sub) => sub,
None => return Ok(()),
};
if sub.name == "list" {
return self.answer_list(ctx, interaction, state).await;
}
let Some(verb) = verb_for_subcommand(sub.name) else {
return Ok(());
};
let selector = selector_for(verb, &sub.value);
self.answer_verb(verb, selector, ctx, interaction, state)
.await
})
}
fn autocomplete<'a>(
&'a self,
ctx: &'a Context,
interaction: &'a CommandInteraction,
state: &'a State,
) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
Box::pin(async move {
let Some(focused) = interaction.data.autocomplete() else {
return;
};
let choices = self.suggestions(&state.live, focused.value).await;
let response = CreateAutocompleteResponse::new().set_choices(
choices
.into_iter()
.map(AutocompleteChoice::from)
.collect::<Vec<_>>(),
);
if let Err(err) = interaction
.create_response(ctx, CreateInteractionResponse::Autocomplete(response))
.await
{
eprintln!("{}", autocomplete_failed_message(&err));
}
})
}
}
#[cfg(test)]
mod tests {
use shep_client::shep_core::protocol::{Request, Response};
use super::*;
use crate::test_support::{dog_sample, sample, test_live};
#[test]
fn the_verbs_that_need_a_name_declare_it_required() {
let json = serde_json::to_value(ShepCommand.data()).expect("json");
let sub: Vec<_> = json["options"]
.as_array()
.expect("options")
.iter()
.collect();
for verb in ["start", "stop", "restart", "reload", "delete", "flush"] {
let option = sub.iter().find(|o| o["name"] == verb).expect(verb);
let name_option = option["options"]
.as_array()
.expect("sub options")
.iter()
.find(|o| o["name"] == "name")
.expect("name");
assert_eq!(name_option["required"], true, "{verb} needs a name");
}
for verb in ["list", "save", "reopen"] {
assert!(sub.iter().any(|o| o["name"] == verb), "{verb} missing");
}
}
#[test]
fn reopen_declares_an_optional_autocompleting_name() {
let json = serde_json::to_value(ShepCommand.data()).expect("json");
let sub: Vec<_> = json["options"]
.as_array()
.expect("options")
.iter()
.collect();
let reopen = sub.iter().find(|o| o["name"] == "reopen").expect("reopen");
let name_option = reopen["options"]
.as_array()
.expect("sub options")
.iter()
.find(|o| o["name"] == "name")
.expect("name");
assert_eq!(
name_option["required"], false,
"reopen's name must not be required"
);
assert_eq!(
name_option["autocomplete"], true,
"reopen's name must autocomplete"
);
for verb in ["list", "save"] {
let option = sub.iter().find(|o| o["name"] == verb).expect(verb);
assert!(
option["options"].as_array().is_none_or(Vec::is_empty),
"{verb} should take no options"
);
}
}
#[test]
fn reopen_with_no_name_selects_everything() {
assert_eq!(selector_for_name(Verb::Reopen, None), SelectorSpec::All);
}
#[test]
fn reopen_with_a_name_selects_that_sheep() {
assert_eq!(
selector_for_name(Verb::Reopen, Some("web".to_owned())),
SelectorSpec::Name("web".to_owned())
);
}
#[test]
fn save_always_selects_everything_regardless_of_a_name() {
assert_eq!(selector_for_name(Verb::Save, None), SelectorSpec::All);
assert_eq!(
selector_for_name(Verb::Save, Some("web".to_owned())),
SelectorSpec::All
);
}
#[test]
fn a_verb_needing_a_name_selects_the_name_it_was_given() {
assert_eq!(
selector_for_name(Verb::Start, Some("web".to_owned())),
SelectorSpec::Name("web".to_owned())
);
assert_eq!(
selector_for_name(Verb::Start, None),
SelectorSpec::Name(String::new())
);
}
#[tokio::test]
async fn autocomplete_offers_the_live_flock() {
let (live, mut fake) = test_live().await;
fake.expect(Request::ListFlock)
.answer(Response::Flock(vec![sample("web"), sample("api")]));
let offered = ShepCommand.suggestions(&live, "").await;
assert_eq!(offered, vec!["web", "api"]);
}
#[tokio::test]
async fn autocomplete_answers_empty_rather_than_erroring() {
let (live, mut fake) = test_live().await;
fake.expect(Request::ListFlock)
.answer_wrong_shape("no shepherd");
assert!(ShepCommand.suggestions(&live, "").await.is_empty());
}
#[tokio::test]
async fn ignore_dogs_hides_dogs_from_the_listing() {
let (live, mut fake) = test_live().await;
fake.expect(Request::ListFlock)
.answer(Response::Flock(vec![sample("web"), dog_sample("bark")]));
let listed = ShepCommand
.flock_for_listing(&live, true)
.await
.expect("ok");
assert_eq!(listed.len(), 1);
assert_eq!(listed[0].name, "web");
}
#[tokio::test]
async fn autocomplete_filters_by_the_typed_partial() {
let (live, mut fake) = test_live().await;
fake.expect(Request::ListFlock).answer(Response::Flock(vec![
sample("web"),
sample("worker"),
sample("api"),
]));
let offered = ShepCommand.suggestions(&live, "WO").await;
assert_eq!(offered, vec!["worker"]);
}
#[tokio::test]
async fn a_name_past_the_choice_limit_is_dropped_not_offered() {
let (live, mut fake) = test_live().await;
let huge = "x".repeat(limits::AUTOCOMPLETE_CHOICE_NAME_LIMIT + 1);
fake.expect(Request::ListFlock)
.answer(Response::Flock(vec![sample("web"), sample(&huge)]));
let offered = ShepCommand.suggestions(&live, "").await;
assert_eq!(offered, vec!["web"]);
}
#[test]
fn every_declared_subcommand_but_list_names_a_verb() {
let json = serde_json::to_value(ShepCommand.data()).expect("json");
for option in json["options"].as_array().expect("options") {
let name = option["name"].as_str().expect("name");
if name == "list" {
assert_eq!(verb_for_subcommand(name), None);
} else {
assert!(verb_for_subcommand(name).is_some(), "{name} names no Verb");
}
}
}
#[test]
fn a_reply_past_the_content_limit_is_fitted_not_refused() {
let huge = "x".repeat(limits::MESSAGE_CONTENT_LIMIT * 2);
let fitted = reply_content(&huge);
assert_eq!(fitted.chars().count(), limits::MESSAGE_CONTENT_LIMIT);
}
#[test]
fn a_short_reply_is_unchanged() {
assert_eq!(reply_content("started web"), "started web");
}
#[test]
fn thirty_sheep_exceed_the_embed_count_cap() {
let flock: Vec<ProcessInfo> = (0..30).map(sample_numbered).collect();
let groups = pack_by_budget(
flock,
limits::MESSAGE_CHARACTER_BUDGET,
limits::EMBED_MAX_COUNT,
embed::embed_character_count,
);
assert!(groups.len() > 1, "thirty sheep landed on one message");
let mut seen = 0;
for group in &groups {
assert!(group.len() <= limits::EMBED_MAX_COUNT, "{}", group.len());
let total: usize = group.iter().map(embed::embed_character_count).sum();
assert!(
total <= limits::MESSAGE_CHARACTER_BUDGET,
"{total} over the message budget"
);
seen += group.len();
}
assert_eq!(seen, 30, "a sheep was lost or duplicated while packing");
}
#[test]
fn two_worst_case_sheep_exceed_the_character_budget() {
let flock = vec![
crate::test_support::worst_case_sample(1),
crate::test_support::worst_case_sample(2),
];
let groups = pack_by_budget(
flock,
limits::MESSAGE_CHARACTER_BUDGET,
limits::EMBED_MAX_COUNT,
embed::embed_character_count,
);
assert_eq!(groups.len(), 2, "two worst-case sheep shared one message");
let mut seen = 0;
for group in &groups {
assert_eq!(group.len(), 1, "a worst-case sheep shared its message");
let total: usize = group.iter().map(embed::embed_character_count).sum();
assert!(
total <= limits::MESSAGE_CHARACTER_BUDGET,
"{total} over the message budget"
);
seen += group.len();
}
assert_eq!(seen, 2, "a sheep was lost or duplicated while packing");
}
fn sample_numbered(id: u32) -> ProcessInfo {
ProcessInfo::builder(
id,
format!("sheep-{id}"),
shep_client::shep_core::status::ProcStatus::Online,
)
.build()
}
#[test]
fn nothing_printed_for_a_person_carries_a_dash() {
crate::test_support::assert_no_dashes(empty_flock_message());
crate::test_support::assert_no_dashes(&autocomplete_failed_message(
&serenity::Error::Other("refused"),
));
let json = serde_json::to_value(ShepCommand.data()).expect("json");
crate::test_support::assert_no_dashes_deep(&json);
}
}