test: update generate image in chat test (#1145)

tests/ai_test/chat_test.rs

@@ -1,4 +1,5 @@
-use crate::ai_test::util::read_text_from_asset;
+use crate::ai_test::util::{extract_image_url, read_text_from_asset};
+use std::time::Duration;
 
 use appflowy_ai_client::dto::{
   ChatQuestionQuery, OutputContent, OutputContentMetadata, OutputLayout, ResponseFormat,
@@ -360,7 +361,7 @@ async fn get_text_with_image_message_test() {
     .unwrap();
 
   let query = ChatQuestionQuery {
-    chat_id,
+    chat_id: chat_id.clone(),
     question_id: question.message_id,
     format: ResponseFormat {
       output_layout: OutputLayout::SimpleTable,
@@ -381,6 +382,45 @@ async fn get_text_with_image_message_test() {
     .unwrap();
   let answer = collect_answer(answer_stream).await;
   println!("answer:\n{}", answer);
+  let image_url = extract_image_url(&answer).unwrap();
+  let (workspace_id_url, chat_id_url, file_id_url) = test_client
+    .api_client
+    .parse_blob_url_v1(&image_url)
+    .unwrap();
+  assert_eq!(workspace_id, workspace_id_url);
+  assert_eq!(chat_id, chat_id_url);
+
+  let mut retries = 3;
+  let retry_interval = Duration::from_secs(8);
+  let mut last_error = None;
+
+  // The image will be generated in the background, so we need to retry until it's available
+  while retries > 0 {
+    match test_client
+      .api_client
+      .get_blob_v1(&workspace_id_url, &chat_id_url, &file_id_url)
+      .await
+    {
+      Ok(_) => {
+        // Success, exit the loop
+        last_error = None;
+        break;
+      },
+      Err(err) => {
+        // Save the error and retry
+        last_error = Some(err);
+        retries -= 1;
+        if retries > 0 {
+          tokio::time::sleep(retry_interval).await;
+        }
+      },
+    }
+  }
+
+  if let Some(err) = last_error {
+    panic!("Failed to get blob after retries: {:?}", err);
+  }
+
   assert!(!answer.is_empty());
 }
 

tests/ai_test/util.rs

@@ -1,3 +1,4 @@
+use fancy_regex::Regex;
 use std::fs::File;
 use std::io::Read;
 
@@ -7,3 +8,16 @@ pub(crate) fn read_text_from_asset(file_name: &str) -> String {
   file.read_to_end(&mut buffer).unwrap();
   String::from_utf8(buffer).unwrap()
 }
+
+pub fn extract_image_url(text: &str) -> Option<String> {
+  // Define a regex pattern to match the image URL inside ![]()
+  let re = Regex::new(r"!\[\]\((https?://[^\s]+)\)").unwrap();
+
+  // Search for the first match in the text
+  if let Ok(Some(captures)) = re.captures(text) {
+    // Extract the matched group (the URL)
+    captures.get(1).map(|m| m.as_str().to_string())
+  } else {
+    None
+  }
+}