AI Add 'sbom-info' tool test

Signed-off-by: mrizzi <[email protected]>
mrizzi · Nov 13, 2024 · bb742b1 · bb742b1
1 parent a6e9a09
commit bb742b1
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 6 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/modules/fundamental/Cargo.toml b/modules/fundamental/Cargo.toml
@@ -68,8 +68,9 @@ tokio-util = { workspace = true }
 trustify-cvss = { workspace = true }
 trustify-test-context = { workspace = true }
 urlencoding = { workspace = true }
-walkdir = { workspace =  true }
+walkdir = { workspace = true }
 zip = { workspace = true }
+termimad = "0.31.0"
 
 [[bench]]
 name = "bench"

diff --git a/modules/fundamental/src/ai/endpoints/test.rs b/modules/fundamental/src/ai/endpoints/test.rs
@@ -23,7 +23,7 @@ async fn configure(ctx: &TrustifyContext) -> anyhow::Result<()> {
 
     let app = caller(ctx).await?;
     let mut req = ChatState::new();
-    req.add_human_message("What is the latest version of Trusted Profile Analyzer?".into());
+    req.add_human_message("Give me information about the SBOMs available for quarkus reporting its name, SHA and URL.".into());
 
     let request = TestRequest::post()
         .uri("/api/v1/ai/completions")
@@ -36,7 +36,12 @@ async fn configure(ctx: &TrustifyContext) -> anyhow::Result<()> {
 
     let result: ChatState = actix_web::test::read_body_json(response).await;
     log::info!("result: {:?}", result);
-    assert!(result.messages.last().unwrap().content.contains("37.17.9"));
+    assert!(result
+        .messages
+        .last()
+        .unwrap()
+        .content
+        .contains("quarkus-bom"));
 
     Ok(())
 }

diff --git a/modules/fundamental/src/ai/service/test.rs b/modules/fundamental/src/ai/service/test.rs
@@ -37,6 +37,8 @@ pub async fn ingest_fixtures(ctx: &TrustifyContext) -> Result<(), anyhow::Error>
 
     ctx.ingest_documents(["osv/RUSTSEC-2021-0079.json", "cve/CVE-2021-32714.json"])
         .await?;
+    ctx.ingest_document("quarkus/v1/quarkus-bom-2.13.8.Final-redhat-00004.json")
+        .await?;
 
     Ok(())
 }
@@ -70,12 +72,22 @@ async fn completions(ctx: &TrustifyContext) -> Result<(), anyhow::Error> {
     ingest_fixtures(ctx).await?;
 
     let mut req = ChatState::new();
-    req.add_human_message("What is the latest version of Trusted Profile Analyzer?".into());
+    req.add_human_message(
+        "Give me information about the SBOMs available for quarkus reporting its name, SHA and URL."
+            .into(),
+    );
 
     let result = service.completions(&req, ()).await?;
 
-    log::info!("result: {:?}", result);
-    assert!(result.messages.last().unwrap().content.contains("37.17.9"));
+    log::info!("result: {:#?}", result);
+    let last_message_content = result.messages.last().unwrap().content.clone();
+    println!(
+        "Test formatted output:\n\n{}\n",
+        termimad::inline(last_message_content.as_str())
+    );
+    assert!(last_message_content.contains("quarkus-bom"));
+    assert!(last_message_content
+        .contains("5a370574a991aa42f7ecc5b7d88754b258f81c230a73bea247c0a6fcc6f608ab"));
 
     Ok(())
 }