From b6fd884387a78b37026c60764ad00286888e98a7 Mon Sep 17 00:00:00 2001
From: "D. Scott Boggs" <scott@tams.tech>
Date: Wed, 25 Jan 2023 06:59:43 -0500
Subject: [PATCH] Add client method to wait for attachment processing

---
 Cargo.toml                 |  1 +
 entities/src/attachment.rs | 24 +++++++++++++++++++
 src/lib.rs                 |  2 ++
 src/mastodon.rs            | 48 ++++++++++++++++++++++++++++++++++++--
 src/polling_time.rs        | 18 ++++++++++++++
 5 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 src/polling_time.rs

diff --git a/Cargo.toml b/Cargo.toml
index 747bee4..81bede5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -44,6 +44,7 @@ futures-util = "0.3.25"
 static_assertions = "1.1.0"
 percent-encoding = "2.2.0"
 thiserror = "1.0.38"
+derive_deref = "1.1.1"
 
 [dependencies.parse_link_header]
 version = "0.3.3"
diff --git a/entities/src/attachment.rs b/entities/src/attachment.rs
index 15fc2a4..32c42a2 100644
--- a/entities/src/attachment.rs
+++ b/entities/src/attachment.rs
@@ -26,6 +26,7 @@ pub struct Attachment {
     /// Noop will be removed.
     pub description: Option<String>,
 }
+
 /// Wrapper type for a attachment ID string
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 #[serde(transparent)]
@@ -99,3 +100,26 @@ pub enum MediaType {
     #[serde(rename = "unknown")]
     Unknown,
 }
+
+/// A media attachment which has been processed and has a URL.
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+pub struct ProcessedAttachment {
+    /// ID of the attachment.
+    pub id: AttachmentId,
+    /// The media type of an attachment.
+    #[serde(rename = "type")]
+    pub media_type: MediaType,
+    /// URL of the locally hosted version of the image.
+    pub url: String,
+    /// For remote images, the remote URL of the original image.
+    pub remote_url: Option<String>,
+    /// URL of the preview image.
+    pub preview_url: String,
+    /// Shorter URL for the image, for insertion into text
+    /// (only present on local images)
+    pub text_url: Option<String>,
+    /// Meta information about the attachment.
+    pub meta: Option<Meta>,
+    /// Noop will be removed.
+    pub description: Option<String>,
+}
diff --git a/src/lib.rs b/src/lib.rs
index c5a418f..60dc78e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -123,6 +123,8 @@ pub mod scopes;
 pub mod status_builder;
 #[macro_use]
 mod macros;
+/// How much time to wait before checking an endpoint again.
+pub mod polling_time;
 /// Automatically import the things you need
 pub mod prelude {
     pub use crate::{
diff --git a/src/mastodon.rs b/src/mastodon.rs
index 5c6ecb4..f59b4e7 100644
--- a/src/mastodon.rs
+++ b/src/mastodon.rs
@@ -11,11 +11,14 @@ use crate::{
     errors::{Error, Result},
     event_stream::event_stream,
     helpers::read_response::read_response,
-    log_serde, AddFilterRequest, AddPushRequest, Data, NewStatus, Page, StatusesRequest,
-    UpdateCredsRequest, UpdatePushRequest,
+    log_serde,
+    polling_time::PollingTime,
+    AddFilterRequest, AddPushRequest, Data, NewStatus, Page, StatusesRequest, UpdateCredsRequest,
+    UpdatePushRequest,
 };
 use futures::TryStream;
 use log::{as_debug, as_serde, debug, error, trace};
+use mastodon_async_entities::attachment::ProcessedAttachment;
 use reqwest::{multipart::Part, Client, RequestBuilder};
 use url::Url;
 use uuid::Uuid;
@@ -119,6 +122,7 @@ impl Mastodon {
         (delete) delete_from_suggestions[AccountId]: "suggestions/{}" => Empty,
         (post) endorse_user[AccountId]: "accounts/{}/pin" => Relationship,
         (post) unendorse_user[AccountId]: "accounts/{}/unpin" => Relationship,
+        (get) attachment[AttachmentId]: "media/{}" => Attachment,
     }
 
     streaming! {
@@ -326,6 +330,46 @@ impl Mastodon {
         self.following(&me.id).await
     }
 
+    /// Wait for the media to be done processing and return it with the URL.
+    ///
+    /// `Default::default()` may be passed as the polling time to select a
+    /// polling time of 500ms.
+    ///
+    /// ## Example
+    /// ```rust,no_run
+    /// use mastodon_async::prelude::*;
+    /// let mastodon = Mastodon::from(Data::default());
+    /// tokio_test::block_on(async {
+    ///     let attachment = mastodon.media("/path/to/some/file.jpg", None).await.expect("upload");
+    ///     let attachment = mastodon.wait_for_processing(attachment, Default::default()).await.expect("processing");
+    ///     println!("{}", attachment.url);
+    /// });
+    /// ```
+    pub async fn wait_for_processing(
+        &self,
+        mut attachment: Attachment,
+        polling_time: PollingTime,
+    ) -> Result<ProcessedAttachment> {
+        let id = attachment.id;
+        loop {
+            if let Some(url) = attachment.url {
+                return Ok(ProcessedAttachment {
+                    id,
+                    media_type: attachment.media_type,
+                    url,
+                    remote_url: attachment.remote_url,
+                    preview_url: attachment.preview_url,
+                    text_url: attachment.text_url,
+                    meta: attachment.meta,
+                    description: attachment.description,
+                });
+            } else {
+                attachment = self.attachment(&id).await?;
+                tokio::time::sleep(*polling_time).await;
+            }
+        }
+    }
+
     /// Set the bearer authentication token
     fn authenticated(&self, request: RequestBuilder) -> RequestBuilder {
         request.bearer_auth(&self.data.token)
diff --git a/src/polling_time.rs b/src/polling_time.rs
new file mode 100644
index 0000000..1b80a5f
--- /dev/null
+++ b/src/polling_time.rs
@@ -0,0 +1,18 @@
+use derive_deref::Deref;
+use std::time::Duration;
+
+/// How long to wait before checking an endpoint again.
+#[derive(Debug, Deref, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub struct PollingTime(Duration);
+
+impl Default for PollingTime {
+    fn default() -> Self {
+        Self(Duration::from_millis(500))
+    }
+}
+
+impl From<Duration> for PollingTime {
+    fn from(value: Duration) -> Self {
+        Self(value)
+    }
+}