1
Fork 0

Set up testing infrastructure

This commit is contained in:
prescientmoon 2024-09-06 17:31:20 +02:00
parent e74ddfd106
commit cba88c5def
Signed by: prescientmoon
SSH key fingerprint: SHA256:WFp/cO76nbarETAoQcQXuV+0h7XJsEsOCI0UsyPIy6U
18 changed files with 494 additions and 176 deletions

View file

@ -110,12 +110,12 @@ async fn info(
/// Show the best score on a given chart
#[poise::command(prefix_command, slash_command, user_cooldown = 1)]
async fn best(
ctx: Context<'_>,
mut ctx: Context<'_>,
#[rest]
#[description = "Name of chart to show (difficulty at the end)"]
name: String,
) -> Result<(), Error> {
let user = get_user!(&ctx);
let user = get_user!(&mut ctx);
let (song, chart) = guess_song_and_chart(&ctx.data(), &name)?;
let play = ctx
@ -164,13 +164,13 @@ async fn best(
/// Show the best score on a given chart
#[poise::command(prefix_command, slash_command, user_cooldown = 10)]
async fn plot(
ctx: Context<'_>,
mut ctx: Context<'_>,
scoring_system: Option<ScoringSystem>,
#[rest]
#[description = "Name of chart to show (difficulty at the end)"]
name: String,
) -> Result<(), Error> {
let user = get_user!(&ctx);
let user = get_user!(&mut ctx);
let scoring_system = scoring_system.unwrap_or_default();
let (song, chart) = guess_song_and_chart(&ctx.data(), &name)?;

210
src/commands/discord.rs Normal file
View file

@ -0,0 +1,210 @@
use std::num::NonZeroU64;
use poise::serenity_prelude::{futures::future::join_all, CreateAttachment, CreateMessage};
use crate::{
context::{Error, UserContext},
timed,
};
// {{{ Trait
pub trait MessageContext {
/// Get the user context held by the message
fn data(&self) -> &UserContext;
fn author_id(&self) -> u64;
/// Reply to the current message
async fn reply(&mut self, text: &str) -> Result<(), Error>;
/// Deliver a message containing references to files.
async fn send_files(
&mut self,
attachments: impl IntoIterator<Item = CreateAttachment>,
message: CreateMessage,
) -> Result<(), Error>;
// {{{ Input attachments
type Attachment;
fn is_image(attachment: &Self::Attachment) -> bool;
fn filename(attachment: &Self::Attachment) -> &str;
fn attachment_id(attachment: &Self::Attachment) -> NonZeroU64;
/// Downloads a single file
async fn download(&self, attachment: &Self::Attachment) -> Result<Vec<u8>, Error>;
/// Downloads every image
async fn download_images<'a>(
&self,
attachments: &'a Vec<Self::Attachment>,
) -> Result<Vec<(&'a Self::Attachment, Vec<u8>)>, Error> {
let download_tasks = attachments
.iter()
.filter(|file| Self::is_image(file))
.map(|file| async move { (file, self.download(file).await) });
let downloaded = timed!("dowload_files", { join_all(download_tasks).await });
downloaded
.into_iter()
.map(|(file, bytes)| Ok((file, bytes?)))
.collect::<Result<_, Error>>()
}
// }}}
}
// }}}
// {{{ Poise implementation
impl<'a, 'b> MessageContext
for poise::Context<'a, UserContext, Box<dyn std::error::Error + Send + Sync + 'b>>
{
type Attachment = poise::serenity_prelude::Attachment;
fn data(&self) -> &UserContext {
Self::data(*self)
}
fn author_id(&self) -> u64 {
self.author().id.get()
}
async fn reply(&mut self, text: &str) -> Result<(), Error> {
Self::reply(*self, text).await?;
Ok(())
}
async fn send_files(
&mut self,
attachments: impl IntoIterator<Item = CreateAttachment>,
message: CreateMessage,
) -> Result<(), Error> {
self.channel_id()
.send_files(self.http(), attachments, message)
.await?;
Ok(())
}
// {{{ Input attachments
fn attachment_id(attachment: &Self::Attachment) -> NonZeroU64 {
NonZeroU64::new(attachment.id.get()).unwrap()
}
fn filename(attachment: &Self::Attachment) -> &str {
&attachment.filename
}
fn is_image(attachment: &Self::Attachment) -> bool {
attachment.dimensions().is_some()
}
async fn download(&self, attachment: &Self::Attachment) -> Result<Vec<u8>, Error> {
let res = poise::serenity_prelude::Attachment::download(attachment).await?;
Ok(res)
}
// }}}
}
// }}}
// {{{ Testing context
pub mod mock {
use std::{env, fs, path::PathBuf};
use super::*;
pub struct MockContext {
pub user_id: u64,
pub data: UserContext,
pub messages: Vec<(CreateMessage, Vec<CreateAttachment>)>,
}
impl MockContext {
pub fn new(data: UserContext) -> Self {
Self {
data,
user_id: 666,
messages: vec![],
}
}
pub fn write_to(&self, path: &PathBuf) -> Result<(), Error> {
if env::var("SHIMMERING_TEST_REGEN").unwrap_or_default() == "1" {
fs::remove_dir_all(path)?;
}
fs::create_dir_all(path)?;
for (i, (message, attachments)) in self.messages.iter().enumerate() {
let dir = path.join(format!("{i}"));
fs::create_dir_all(&dir)?;
let message_file = dir.join("message.toml");
if message_file.exists() {
assert_eq!(
toml::to_string_pretty(message)?,
fs::read_to_string(message_file)?
);
} else {
fs::write(&message_file, toml::to_string_pretty(message)?)?;
}
for attachment in attachments {
let path = dir.join(&attachment.filename);
if path.exists() {
assert_eq!(&attachment.data, &fs::read(path)?);
} else {
fs::write(&path, &attachment.data)?;
}
}
}
Ok(())
}
}
impl MessageContext for MockContext {
fn author_id(&self) -> u64 {
self.user_id
}
fn data(&self) -> &UserContext {
&self.data
}
async fn reply(&mut self, text: &str) -> Result<(), Error> {
self.messages
.push((CreateMessage::new().content(text), Vec::new()));
Ok(())
}
async fn send_files(
&mut self,
attachments: impl IntoIterator<Item = CreateAttachment>,
message: CreateMessage,
) -> Result<(), Error> {
self.messages
.push((message, attachments.into_iter().collect()));
Ok(())
}
// {{{ Input attachments
type Attachment = PathBuf;
fn filename(attachment: &Self::Attachment) -> &str {
attachment.file_name().unwrap().to_str().unwrap()
}
// This is a dumb implementation, but it works for testing...
fn is_image(attachment: &Self::Attachment) -> bool {
let ext = attachment.extension().unwrap();
ext == "png" || ext == "jpg" || ext == "webp"
}
fn attachment_id(_attachment: &Self::Attachment) -> NonZeroU64 {
NonZeroU64::new(666).unwrap()
}
async fn download(&self, attachment: &Self::Attachment) -> Result<Vec<u8>, Error> {
let res = tokio::fs::read(attachment).await?;
Ok(res)
}
// }}}
}
}
// }}}

View file

@ -1,6 +1,7 @@
use crate::context::{Context, Error};
pub mod chart;
pub mod discord;
pub mod score;
pub mod stats;
pub mod utils;

View file

@ -1,15 +1,14 @@
use std::time::Instant;
use crate::arcaea::play::{CreatePlay, Play};
use crate::arcaea::score::Score;
use crate::context::{Context, Error};
use crate::recognition::recognize::{ImageAnalyzer, ScoreKind};
use crate::user::{discord_id_to_discord_user, User};
use crate::{edit_reply, get_user, timed};
use crate::{get_user, timed};
use image::DynamicImage;
use poise::serenity_prelude::futures::future::join_all;
use poise::serenity_prelude as serenity;
use poise::serenity_prelude::CreateMessage;
use poise::{serenity_prelude as serenity, CreateReply};
use super::discord::MessageContext;
// {{{ Score
/// Score management
@ -24,13 +23,13 @@ pub async fn score(_ctx: Context<'_>) -> Result<(), Error> {
}
// }}}
// {{{ Score magic
/// Identify scores from attached images.
#[poise::command(prefix_command, slash_command)]
pub async fn magic(
ctx: Context<'_>,
#[description = "Images containing scores"] files: Vec<serenity::Attachment>,
// {{{ Implementation
async fn magic_impl<C: MessageContext>(
ctx: &mut C,
files: Vec<C::Attachment>,
) -> Result<(), Error> {
let user = get_user!(&ctx);
let user = get_user!(ctx);
let files = ctx.download_images(&files).await?;
if files.len() == 0 {
ctx.reply("No images found attached to message").await?;
@ -39,30 +38,9 @@ pub async fn magic(
let mut embeds = Vec::with_capacity(files.len());
let mut attachments = Vec::with_capacity(files.len());
let handle = ctx
.reply(format!("Processed 0/{} scores", files.len()))
.await?;
let mut analyzer = ImageAnalyzer::default();
// {{{ Download files
let download_tasks = files
.iter()
.filter(|file| file.dimensions().is_some())
.map(|file| async move { (file, file.download().await) });
let downloaded = timed!("dowload_files", { join_all(download_tasks).await });
if downloaded.len() < files.len() {
ctx.reply("One or more of the attached files are not images!")
.await?;
}
// }}}
for (i, (file, bytes)) in downloaded.into_iter().enumerate() {
let bytes = bytes?;
let start = Instant::now();
for (i, (attachment, bytes)) in files.into_iter().enumerate() {
// {{{ Preapare image
let mut image = timed!("decode image", { image::load_from_memory(&bytes)? });
let mut grayscale_image = timed!("grayscale image", {
@ -109,7 +87,14 @@ pub async fn magic(
// edit_reply!(ctx, handle, "Image {}: reading score", i + 1).await?;
let score = timed!("read_score", {
analyzer.read_score(ctx.data(), Some(chart.note_count), &grayscale_image, kind)?
analyzer
.read_score(ctx.data(), Some(chart.note_count), &grayscale_image, kind)
.map_err(|err| {
format!(
"Could not read score for chart {} [{:?}]: {err}",
song.title, chart.difficulty
)
})?
});
// {{{ Build play
@ -117,7 +102,7 @@ pub async fn magic(
Score::resolve_distibution_ambiguities(score, note_distribution, chart.note_count);
let play = CreatePlay::new(score)
.with_attachment(file)
.with_attachment(C::attachment_id(attachment))
.with_fars(maybe_fars)
.with_max_recall(max_recall)
.save(&ctx.data(), &user, &chart)?;
@ -136,41 +121,113 @@ pub async fn magic(
if let Err(err) = result {
analyzer
.send_discord_error(ctx, &image, &file.filename, err)
.send_discord_error(ctx, &image, C::filename(&attachment), err)
.await?;
}
let took = start.elapsed();
edit_reply!(
ctx,
handle,
"Processed {}/{} scores. Last score took {took:?} to process.",
i + 1,
files.len()
)
.await?;
}
handle.delete(ctx).await?;
if embeds.len() > 0 {
ctx.channel_id()
.send_files(ctx.http(), attachments, CreateMessage::new().embeds(embeds))
ctx.send_files(attachments, CreateMessage::new().embeds(embeds))
.await?;
}
Ok(())
}
// }}}
// {{{ Tests
#[cfg(test)]
mod magic_tests {
use std::{path::PathBuf, process::Command, str::FromStr};
use r2d2_sqlite::SqliteConnectionManager;
use crate::{
commands::discord::mock::MockContext,
context::{connect_db, get_shared_context},
};
use super::*;
macro_rules! with_ctx {
($test_path:expr, $f:expr) => {{
let mut data = (*get_shared_context().await).clone();
let dir = tempfile::tempdir()?;
let path = dir.path().join("db.sqlite");
println!("path {path:?}");
data.db = connect_db(SqliteConnectionManager::file(path));
Command::new("scripts/import-charts.py")
.env("SHIMMERING_DATA_DIR", dir.path().to_str().unwrap())
.output()
.unwrap();
let mut ctx = MockContext::new(data);
User::create_from_context(&ctx)?;
let res: Result<(), Error> = $f(&mut ctx).await;
res?;
ctx.write_to(&PathBuf::from_str($test_path)?)?;
Ok(())
}};
}
#[tokio::test]
async fn no_pics() -> Result<(), Error> {
with_ctx!("test/commands/score/magic/no_pics", async |ctx| {
magic_impl(ctx, vec![]).await?;
Ok(())
})
}
#[tokio::test]
async fn basic_pic() -> Result<(), Error> {
with_ctx!("test/commands/score/magic/single_pic", async |ctx| {
magic_impl(
ctx,
vec![PathBuf::from_str("test/screenshots/alter_ego.jpg")?],
)
.await?;
Ok(())
})
}
#[tokio::test]
async fn weird_kerning() -> Result<(), Error> {
with_ctx!("test/commands/score/magic/weird_kerning", async |ctx| {
magic_impl(
ctx,
vec![
PathBuf::from_str("test/screenshots/antithese_74_kerning.jpg")?,
PathBuf::from_str("test/screenshots/genocider_24_kerning.jpg")?,
],
)
.await?;
Ok(())
})
}
}
// }}}
/// Identify scores from attached images.
#[poise::command(prefix_command, slash_command)]
pub async fn magic(
mut ctx: Context<'_>,
#[description = "Images containing scores"] files: Vec<serenity::Attachment>,
) -> Result<(), Error> {
magic_impl(&mut ctx, files).await?;
Ok(())
}
// }}}
// {{{ Score delete
/// Delete scores, given their IDs.
#[poise::command(prefix_command, slash_command)]
pub async fn delete(
ctx: Context<'_>,
mut ctx: Context<'_>,
#[description = "Id of score to delete"] ids: Vec<u32>,
) -> Result<(), Error> {
let user = get_user!(&ctx);
let user = get_user!(&mut ctx);
if ids.len() == 0 {
ctx.reply("Empty ID list provided").await?;

View file

@ -43,7 +43,7 @@ pub async fn stats(_ctx: Context<'_>) -> Result<(), Error> {
// }}}
// {{{ Render best plays
async fn best_plays(
ctx: &Context<'_>,
ctx: &mut Context<'_>,
user: &User,
scoring_system: ScoringSystem,
grid_size: (u32, u32),
@ -403,7 +403,7 @@ async fn best_plays(
ImageBuffer::from_raw(width, height, drawer.canvas.buffer.into_vec()).unwrap(),
);
debug_image_log(&image)?;
debug_image_log(&image);
if image.height() > 4096 {
image = image.resize(4096, 4096, image::imageops::FilterType::Nearest);
@ -426,10 +426,10 @@ async fn best_plays(
// {{{ B30
/// Show the 30 best scores
#[poise::command(prefix_command, slash_command, user_cooldown = 30)]
pub async fn b30(ctx: Context<'_>, scoring_system: Option<ScoringSystem>) -> Result<(), Error> {
let user = get_user!(&ctx);
pub async fn b30(mut ctx: Context<'_>, scoring_system: Option<ScoringSystem>) -> Result<(), Error> {
let user = get_user!(&mut ctx);
best_plays(
&ctx,
&mut ctx,
&user,
scoring_system.unwrap_or_default(),
(5, 6),
@ -440,15 +440,15 @@ pub async fn b30(ctx: Context<'_>, scoring_system: Option<ScoringSystem>) -> Res
#[poise::command(prefix_command, slash_command, hide_in_help, global_cooldown = 5)]
pub async fn bany(
ctx: Context<'_>,
mut ctx: Context<'_>,
scoring_system: Option<ScoringSystem>,
width: u32,
height: u32,
) -> Result<(), Error> {
let user = get_user!(&ctx);
let user = get_user!(&mut ctx);
assert_is_pookie!(ctx, user);
best_plays(
&ctx,
&mut ctx,
&user,
scoring_system.unwrap_or_default(),
(width, height),
@ -460,8 +460,8 @@ pub async fn bany(
// {{{ Meta
/// Show stats about the bot itself.
#[poise::command(prefix_command, slash_command, user_cooldown = 1)]
async fn meta(ctx: Context<'_>) -> Result<(), Error> {
let user = get_user!(&ctx);
async fn meta(mut ctx: Context<'_>) -> Result<(), Error> {
let user = get_user!(&mut ctx);
let conn = ctx.data().db.get()?;
let song_count: usize = conn
.prepare_cached("SELECT count() as count FROM songs")?

View file

@ -35,7 +35,7 @@ macro_rules! reply_errors {
match $value {
Ok(v) => v,
Err(err) => {
$ctx.reply(format!("{err}")).await?;
crate::commands::discord::MessageContext::reply($ctx, &format!("{err}")).await?;
return Ok(());
}
}