From 7bf2febd07a87477367b843bd15451021237b6ab Mon Sep 17 00:00:00 2001 From: theoforger Date: Fri, 20 Sep 2024 11:49:05 -0400 Subject: [PATCH 1/6] Add tests module --- src/lib.rs | 2 ++ src/tests.rs | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+) create mode 100644 src/tests.rs diff --git a/src/lib.rs b/src/lib.rs index fa4dbf0..07c18ff 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,8 @@ use std::path::PathBuf; pub mod api; mod clue; mod json_models; +#[cfg(test)] +mod tests; /// Mastermind - An LLM-powered CLI tool to help you be a better spymaster in Codenames #[derive(Parser)] diff --git a/src/tests.rs b/src/tests.rs new file mode 100644 index 0000000..f61251c --- /dev/null +++ b/src/tests.rs @@ -0,0 +1,62 @@ + + +use super::*; + +#[test] +fn test_api_instance() { + let api_instance = api::Instance::new(); + assert!(api_instance.is_ok()); +} + +#[test] +fn test_read_words_from_file() { + let to_link = read_words_from_file(PathBuf::from("./examples/link.txt")); + assert!(to_link.is_ok()); + let to_avoid = read_words_from_file(PathBuf::from("./examples/avoid.txt")); + assert!(to_avoid.is_ok()); +} + +#[tokio::test] +async fn test_language_models_call() { + let api_instance = api::Instance::new(); + assert!(api_instance.is_ok()); + + let to_link = vec![ + "bond".to_string(), + "sound".to_string(), + "park".to_string(), + "penny".to_string(), + "bee".to_string(), + "tokyo".to_string(), + "walrus".to_string(), + "hospital".to_string(), + "scuba diver".to_string(), + ]; + + let to_avoid = vec![ + "angel".to_string(), + "ski".to_string(), + "captain".to_string(), + "bass".to_string(), + "boil".to_string(), + "casino".to_string(), + "star".to_string(), + "fish".to_string(), + "blind".to_string(), + "day".to_string(), + "tip".to_string(), + "goldilocks".to_string(), + "field".to_string(), + "file".to_string(), + "cotton".to_string(), + "scarecrow".to_string(), + "extra virgin olive oil".to_string(), + ]; + + let models = api_instance + .unwrap() + .fetch_clue_collection(to_link, to_avoid) + .await; + + assert!(models.is_ok()); +} From 2290aa03728a9b10680b6b0d80e474c5ed2d4e40 Mon Sep 17 00:00:00 2001 From: theoforger Date: Fri, 20 Sep 2024 17:35:21 -0400 Subject: [PATCH 2/6] Add 3 test units --- Cargo.toml | 3 +- src/api/chat_completions.rs | 11 +- src/api/language_models.rs | 6 +- src/api/mod.rs | 4 + src/tests.rs | 62 --------- src/tests/mock_responses/language_models.json | 122 ++++++++++++++++++ src/tests/mod.rs | 57 ++++++++ 7 files changed, 192 insertions(+), 73 deletions(-) delete mode 100644 src/tests.rs create mode 100644 src/tests/mock_responses/language_models.json create mode 100644 src/tests/mod.rs diff --git a/Cargo.toml b/Cargo.toml index 191b8b2..66e4003 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,4 +11,5 @@ tokio = { version = "1.40.0", features = ["rt", "rt-multi-thread", "macros"] } dotenv = "0.15.0" serde_json = "1.0.128" serde = { version = "1.0.210", features = ["derive"] } -comfy-table = "7.1.1" \ No newline at end of file +comfy-table = "7.1.1" +httpmock = "0.7.0" \ No newline at end of file diff --git a/src/api/chat_completions.rs b/src/api/chat_completions.rs index 9126d48..759fc9e 100644 --- a/src/api/chat_completions.rs +++ b/src/api/chat_completions.rs @@ -1,6 +1,6 @@ -use crate::json_models::chat_completion::ChatCompletionResponse; use super::Instance; use crate::clue::ClueCollection; +use crate::json_models::chat_completion::ChatCompletionResponse; use serde_json::json; const SYSTEM_PROMPT: &str = r#" @@ -36,21 +36,18 @@ impl Instance { .json(&request_body) .send() .await - .map_err(|_| "Failed to fetch clue collection from API server")?; + .map_err(|e| format!("Failed to fetch clue collection from API server: {}", e))?; let parsed_response = response .json::() .await - .map_err(|_| "Failed to parse clues from API server")?; + .map_err(|e| format!("Failed to parse clues from API server: {}", e))?; // Extract usage information from the parsed response let token_usage = parsed_response.usage; // Extract clue strings from the parsed response - let clue_strings = parsed_response - .choices - .first() - .ok_or("Failed to parse clues from API server")? + let clue_strings = parsed_response.choices[0] .message .content .lines() diff --git a/src/api/language_models.rs b/src/api/language_models.rs index 34c9890..8b579b0 100644 --- a/src/api/language_models.rs +++ b/src/api/language_models.rs @@ -1,5 +1,5 @@ -use crate::json_models::language_model::ModelsResponse; use super::Instance; +use crate::json_models::language_model::ModelsResponse; impl Instance { pub async fn fetch_all_model_ids(&self) -> Result, Box> { @@ -9,12 +9,12 @@ impl Instance { .bearer_auth(&self.key) .send() .await - .map_err(|_| "Failed to fetch model IDs from API server")?; + .map_err(|e| format!("Failed to fetch model IDs from API server: {}", e))?; let mut all_model_ids = response .json::() .await - .map_err(|_| "Failed to parse model IDs from API server")? + .map_err(|e| format!("Failed to parse model IDs from API server: {}", e))? .data .iter() .map(|model| model.id.trim().to_string()) diff --git a/src/api/mod.rs b/src/api/mod.rs index 004cef5..f487870 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -54,4 +54,8 @@ impl Instance { self.model_id = model_id; Ok(()) } + + pub fn set_base_url(&mut self, base_url: String) { + self.base_url = base_url; + } } diff --git a/src/tests.rs b/src/tests.rs deleted file mode 100644 index f61251c..0000000 --- a/src/tests.rs +++ /dev/null @@ -1,62 +0,0 @@ - - -use super::*; - -#[test] -fn test_api_instance() { - let api_instance = api::Instance::new(); - assert!(api_instance.is_ok()); -} - -#[test] -fn test_read_words_from_file() { - let to_link = read_words_from_file(PathBuf::from("./examples/link.txt")); - assert!(to_link.is_ok()); - let to_avoid = read_words_from_file(PathBuf::from("./examples/avoid.txt")); - assert!(to_avoid.is_ok()); -} - -#[tokio::test] -async fn test_language_models_call() { - let api_instance = api::Instance::new(); - assert!(api_instance.is_ok()); - - let to_link = vec![ - "bond".to_string(), - "sound".to_string(), - "park".to_string(), - "penny".to_string(), - "bee".to_string(), - "tokyo".to_string(), - "walrus".to_string(), - "hospital".to_string(), - "scuba diver".to_string(), - ]; - - let to_avoid = vec![ - "angel".to_string(), - "ski".to_string(), - "captain".to_string(), - "bass".to_string(), - "boil".to_string(), - "casino".to_string(), - "star".to_string(), - "fish".to_string(), - "blind".to_string(), - "day".to_string(), - "tip".to_string(), - "goldilocks".to_string(), - "field".to_string(), - "file".to_string(), - "cotton".to_string(), - "scarecrow".to_string(), - "extra virgin olive oil".to_string(), - ]; - - let models = api_instance - .unwrap() - .fetch_clue_collection(to_link, to_avoid) - .await; - - assert!(models.is_ok()); -} diff --git a/src/tests/mock_responses/language_models.json b/src/tests/mock_responses/language_models.json new file mode 100644 index 0000000..8eb52c1 --- /dev/null +++ b/src/tests/mock_responses/language_models.json @@ -0,0 +1,122 @@ +{ + "object": "list", + "data": [ + { + "id": "llava-v1.5-7b-4096-preview", + "object": "model", + "created": 1725402373, + "owned_by": "Other", + "active": true, + "context_window": 4096, + "public_apps": null + }, + { + "id": "gemma-7b-it", + "object": "model", + "created": 1693721698, + "owned_by": "Google", + "active": true, + "context_window": 8192, + "public_apps": null + }, + { + "id": "llama-3.1-8b-instant", + "object": "model", + "created": 1693721698, + "owned_by": "Meta", + "active": true, + "context_window": 131072, + "public_apps": null + }, + { + "id": "whisper-large-v3", + "object": "model", + "created": 1693721698, + "owned_by": "OpenAI", + "active": true, + "context_window": 448, + "public_apps": null + }, + { + "id": "mixtral-8x7b-32768", + "object": "model", + "created": 1693721698, + "owned_by": "Mistral AI", + "active": true, + "context_window": 32768, + "public_apps": null + }, + { + "id": "llama3-8b-8192", + "object": "model", + "created": 1693721698, + "owned_by": "Meta", + "active": true, + "context_window": 8192, + "public_apps": null + }, + { + "id": "distil-whisper-large-v3-en", + "object": "model", + "created": 1693721698, + "owned_by": "Hugging Face", + "active": true, + "context_window": 448, + "public_apps": null + }, + { + "id": "llama-guard-3-8b", + "object": "model", + "created": 1693721698, + "owned_by": "Meta", + "active": true, + "context_window": 8192, + "public_apps": null + }, + { + "id": "gemma2-9b-it", + "object": "model", + "created": 1693721698, + "owned_by": "Google", + "active": true, + "context_window": 8192, + "public_apps": null + }, + { + "id": "llama3-70b-8192", + "object": "model", + "created": 1693721698, + "owned_by": "Meta", + "active": true, + "context_window": 8192, + "public_apps": null + }, + { + "id": "llama3-groq-70b-8192-tool-use-preview", + "object": "model", + "created": 1693721698, + "owned_by": "Groq", + "active": true, + "context_window": 8192, + "public_apps": null + }, + { + "id": "llama-3.1-70b-versatile", + "object": "model", + "created": 1693721698, + "owned_by": "Meta", + "active": true, + "context_window": 131072, + "public_apps": null + }, + { + "id": "llama3-groq-8b-8192-tool-use-preview", + "object": "model", + "created": 1693721698, + "owned_by": "Groq", + "active": true, + "context_window": 8192, + "public_apps": null + } + ] +} \ No newline at end of file diff --git a/src/tests/mod.rs b/src/tests/mod.rs new file mode 100644 index 0000000..4e474d9 --- /dev/null +++ b/src/tests/mod.rs @@ -0,0 +1,57 @@ +use super::*; +use crate::api::Instance; +use httpmock::prelude::*; + +#[test] +fn test_api_instance() { + let api_instance = api::Instance::new(); + assert!(api_instance.is_ok()); +} + +#[test] +fn test_read_words_from_file() { + let to_link = read_words_from_file(PathBuf::from("examples/link.txt")); + assert!(to_link.is_ok()); + let to_avoid = read_words_from_file(PathBuf::from("examples/avoid.txt")); + assert!(to_avoid.is_ok()); +} + +#[tokio::test] +async fn test_fetch_language_models() { + // Start a lightweight mock server. + let server = MockServer::start_async().await; + + // Create a mock on the server. + let mock = server.mock(|when, then| { + when.method(GET).path("/models"); + then.status(200) + .header("content-type", "application/json") + .body_from_file("src/tests/mock_responses/language_models.json"); + }); + + // Create an API instance and set the base url to mock server url + let mut api_instance = Instance::new().unwrap(); + api_instance.set_base_url(server.url("/")); + + // Get response from mock server + let response = api_instance.fetch_all_model_ids().await.unwrap(); + mock.assert(); + + // Compare results + let expected_response = [ + "distil-whisper-large-v3-en", + "gemma-7b-it", + "gemma2-9b-it", + "llama-3.1-70b-versatile", + "llama-3.1-8b-instant", + "llama-guard-3-8b", + "llama3-70b-8192", + "llama3-8b-8192", + "llama3-groq-70b-8192-tool-use-preview", + "llama3-groq-8b-8192-tool-use-preview", + "llava-v1.5-7b-4096-preview", + "mixtral-8x7b-32768", + "whisper-large-v3", + ]; + assert_eq!(response, expected_response); +} From 95fd01d2e87cf19fe5dfe958aeef3b1960de74d5 Mon Sep 17 00:00:00 2001 From: theoforger Date: Fri, 20 Sep 2024 17:37:07 -0400 Subject: [PATCH 3/6] Rename `fetch_all_model_ids` to `fetch_language_model_ids` --- src/api/language_models.rs | 2 +- src/api/mod.rs | 2 +- src/main.rs | 2 +- src/tests/mod.rs | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/api/language_models.rs b/src/api/language_models.rs index 8b579b0..a1a00c5 100644 --- a/src/api/language_models.rs +++ b/src/api/language_models.rs @@ -2,7 +2,7 @@ use super::Instance; use crate::json_models::language_model::ModelsResponse; impl Instance { - pub async fn fetch_all_model_ids(&self) -> Result, Box> { + pub async fn fetch_language_model_ids(&self) -> Result, Box> { let response = self .client .get(format!("{}models", self.base_url)) diff --git a/src/api/mod.rs b/src/api/mod.rs index f487870..cec4b7c 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -42,7 +42,7 @@ impl Instance { model_id: String, ) -> Result<(), Box> { // Return Error if the chosen model is not valid - let valid_model_ids = self.fetch_all_model_ids().await?; + let valid_model_ids = self.fetch_language_model_ids().await?; if !valid_model_ids.contains(&model_id) { return Err(format!( "{} is not a valid language model from your provider", diff --git a/src/main.rs b/src/main.rs index e4fdaa5..2089dd4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,7 +13,7 @@ async fn main() -> Result<(), Box> { // If -g is set, call the models API endpoint instead if args.get { - println!("{}", api_instance.fetch_all_model_ids().await?.join("\n")); + println!("{}", api_instance.fetch_language_model_ids().await?.join("\n")); return Ok(()); } diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 4e474d9..0e4181c 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -34,7 +34,7 @@ async fn test_fetch_language_models() { api_instance.set_base_url(server.url("/")); // Get response from mock server - let response = api_instance.fetch_all_model_ids().await.unwrap(); + let response = api_instance.fetch_language_model_ids().await.unwrap(); mock.assert(); // Compare results From 7b0aca106ef3976e5408661d08c1ad7ff6c23b38 Mon Sep 17 00:00:00 2001 From: theoforger Date: Fri, 20 Sep 2024 18:34:18 -0400 Subject: [PATCH 4/6] Add testing for `fetch_clue_collection` --- .../mock_responses/chat_completions.json | 30 +++++++++++++ src/tests/mod.rs | 42 +++++++++++++++++++ 2 files changed, 72 insertions(+) create mode 100644 src/tests/mock_responses/chat_completions.json diff --git a/src/tests/mock_responses/chat_completions.json b/src/tests/mock_responses/chat_completions.json new file mode 100644 index 0000000..008fab6 --- /dev/null +++ b/src/tests/mock_responses/chat_completions.json @@ -0,0 +1,30 @@ +{ + "id": "chatcmpl-869ede85-2f46-4834-a039-28d757e958a5", + "object": "chat.completion", + "created": 1726870549, + "model": "llama-3.1-70b-versatile", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "music, 2, sound, bee\nfilm, 2, bond, tokyo\nfree, 2, park, penny\ndive, 2, scuba diver, hospital\nlarge, 2, walrus, scuba diver" + }, + "logprobs": null, + "finish_reason": "stop" + } + ], + "usage": { + "queue_time": 0.005669406999999987, + "prompt_tokens": 222, + "prompt_time": 0.068204384, + "completion_tokens": 53, + "completion_time": 0.214023764, + "total_tokens": 275, + "total_time": 0.282228148 + }, + "system_fingerprint": "fp_b6828be2c9", + "x_groq": { + "id": "req_01j88r2wfmecr9zgpjn2zmnprb" + } +} \ No newline at end of file diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 0e4181c..ec73ce5 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -55,3 +55,45 @@ async fn test_fetch_language_models() { ]; assert_eq!(response, expected_response); } + +#[tokio::test] +async fn test_fetch_clue_collection() { + // Start a lightweight mock server. + let server = MockServer::start_async().await; + + // Create a mock on the server. + let mock = server.mock(|when, then| { + when.method(POST).path("/chat/completions"); + then.status(200) + .header("content-type", "application/json") + .body_from_file("src/tests/mock_responses/chat_completions.json"); + }); + + // Create an API instance and set the base url to mock server url + let mut api_instance = Instance::new().unwrap(); + api_instance.set_base_url(server.url("/")); + + // Get response from mock server + let response = api_instance + .fetch_clue_collection(vec![], vec![]) + .await + .unwrap(); + mock.assert(); + + // Compare results + let expected_response = " +╭───────┬───────┬───────────────────────╮ +│ Clue ┆ Count ┆ Linked Words │ +╞═══════╪═══════╪═══════════════════════╡ +│ music ┆ 2 ┆ sound, bee │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ film ┆ 2 ┆ bond, tokyo │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ free ┆ 2 ┆ park, penny │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ dive ┆ 2 ┆ scuba diver, hospital │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ large ┆ 2 ┆ walrus, scuba diver │ +╰───────┴───────┴───────────────────────╯"; + assert_eq!(response.generate_table(), expected_response); +} From b2854b283c8ea7975ec579a169b8511986390628 Mon Sep 17 00:00:00 2001 From: theoforger Date: Fri, 20 Sep 2024 18:54:23 -0400 Subject: [PATCH 5/6] Move test results to separate directory --- .../expected_outputs/chat_completions.txt | 13 ++++++ .../expected_outputs/language_models.txt | 13 ++++++ src/tests/mod.rs | 42 ++++--------------- 3 files changed, 35 insertions(+), 33 deletions(-) create mode 100644 src/tests/expected_outputs/chat_completions.txt create mode 100644 src/tests/expected_outputs/language_models.txt diff --git a/src/tests/expected_outputs/chat_completions.txt b/src/tests/expected_outputs/chat_completions.txt new file mode 100644 index 0000000..4106b35 --- /dev/null +++ b/src/tests/expected_outputs/chat_completions.txt @@ -0,0 +1,13 @@ +╭───────┬───────┬───────────────────────╮ +│ Clue ┆ Count ┆ Linked Words │ +╞═══════╪═══════╪═══════════════════════╡ +│ music ┆ 2 ┆ sound, bee │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ film ┆ 2 ┆ bond, tokyo │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ free ┆ 2 ┆ park, penny │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ dive ┆ 2 ┆ scuba diver, hospital │ +├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ +│ large ┆ 2 ┆ walrus, scuba diver │ +╰───────┴───────┴───────────────────────╯ \ No newline at end of file diff --git a/src/tests/expected_outputs/language_models.txt b/src/tests/expected_outputs/language_models.txt new file mode 100644 index 0000000..51b43bf --- /dev/null +++ b/src/tests/expected_outputs/language_models.txt @@ -0,0 +1,13 @@ +distil-whisper-large-v3-en +gemma-7b-it +gemma2-9b-it +llama-3.1-70b-versatile +llama-3.1-8b-instant +llama-guard-3-8b +llama3-70b-8192 +llama3-8b-8192 +llama3-groq-70b-8192-tool-use-preview +llama3-groq-8b-8192-tool-use-preview +llava-v1.5-7b-4096-preview +mixtral-8x7b-32768 +whisper-large-v3 \ No newline at end of file diff --git a/src/tests/mod.rs b/src/tests/mod.rs index ec73ce5..d1cfa91 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -37,23 +37,10 @@ async fn test_fetch_language_models() { let response = api_instance.fetch_language_model_ids().await.unwrap(); mock.assert(); - // Compare results - let expected_response = [ - "distil-whisper-large-v3-en", - "gemma-7b-it", - "gemma2-9b-it", - "llama-3.1-70b-versatile", - "llama-3.1-8b-instant", - "llama-guard-3-8b", - "llama3-70b-8192", - "llama3-8b-8192", - "llama3-groq-70b-8192-tool-use-preview", - "llama3-groq-8b-8192-tool-use-preview", - "llava-v1.5-7b-4096-preview", - "mixtral-8x7b-32768", - "whisper-large-v3", - ]; - assert_eq!(response, expected_response); + // Compare outputs + let output = response.join("\n"); + let expected_output = fs::read_to_string("src/tests/expected_outputs/language_models.txt").unwrap(); + assert_eq!(output, expected_output); } #[tokio::test] @@ -80,20 +67,9 @@ async fn test_fetch_clue_collection() { .unwrap(); mock.assert(); - // Compare results - let expected_response = " -╭───────┬───────┬───────────────────────╮ -│ Clue ┆ Count ┆ Linked Words │ -╞═══════╪═══════╪═══════════════════════╡ -│ music ┆ 2 ┆ sound, bee │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ film ┆ 2 ┆ bond, tokyo │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ free ┆ 2 ┆ park, penny │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ dive ┆ 2 ┆ scuba diver, hospital │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ large ┆ 2 ┆ walrus, scuba diver │ -╰───────┴───────┴───────────────────────╯"; - assert_eq!(response.generate_table(), expected_response); + // Compare outputs + let output = response.generate_table(); + let expected_output = + fs::read_to_string("src/tests/expected_outputs/chat_completions.txt").unwrap(); + assert_eq!(output.trim(), expected_output); } From bec161ba815579460c385ff81130a8efdd3e9434 Mon Sep 17 00:00:00 2001 From: theoforger Date: Fri, 20 Sep 2024 19:18:56 -0400 Subject: [PATCH 6/6] Add function `generate_raw_list` for more consistent testing --- src/clue.rs | 20 +++++++++++++++---- .../expected_outputs/chat_completions.txt | 18 +++++------------ src/tests/mod.rs | 4 ++-- 3 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/clue.rs b/src/clue.rs index eb8dc9d..e5f9392 100644 --- a/src/clue.rs +++ b/src/clue.rs @@ -99,7 +99,7 @@ impl ClueCollection { pub fn display_table(&self) { println!("{}", self.generate_table()); } - + pub fn display_token_info(&self) { eprintln!( "\nToken Usage:\n\ @@ -108,9 +108,21 @@ impl ClueCollection { Completion Tokens: {}\n\ ----------------------\n\ Total Tokens: {}", - self.usage.prompt_tokens, - self.usage.completion_tokens, - self.usage.total_tokens + self.usage.prompt_tokens, self.usage.completion_tokens, self.usage.total_tokens ); } + + pub fn generate_raw_list(&self) -> String { + let mut raw_list = String::new(); + for clue in &self.clues { + let clue_string = format!( + "{} {} - {}\n", + clue.clue_word, + clue.count, + clue.linked_words.join(", ") + ); + raw_list.push_str(clue_string.as_str()); + } + raw_list + } } diff --git a/src/tests/expected_outputs/chat_completions.txt b/src/tests/expected_outputs/chat_completions.txt index 4106b35..6e7afb0 100644 --- a/src/tests/expected_outputs/chat_completions.txt +++ b/src/tests/expected_outputs/chat_completions.txt @@ -1,13 +1,5 @@ -╭───────┬───────┬───────────────────────╮ -│ Clue ┆ Count ┆ Linked Words │ -╞═══════╪═══════╪═══════════════════════╡ -│ music ┆ 2 ┆ sound, bee │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ film ┆ 2 ┆ bond, tokyo │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ free ┆ 2 ┆ park, penny │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ dive ┆ 2 ┆ scuba diver, hospital │ -├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤ -│ large ┆ 2 ┆ walrus, scuba diver │ -╰───────┴───────┴───────────────────────╯ \ No newline at end of file +music 2 - sound, bee +film 2 - bond, tokyo +free 2 - park, penny +dive 2 - scuba diver, hospital +large 2 - walrus, scuba diver diff --git a/src/tests/mod.rs b/src/tests/mod.rs index d1cfa91..5cbbdf3 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -68,8 +68,8 @@ async fn test_fetch_clue_collection() { mock.assert(); // Compare outputs - let output = response.generate_table(); + let output = response.generate_raw_list(); let expected_output = fs::read_to_string("src/tests/expected_outputs/chat_completions.txt").unwrap(); - assert_eq!(output.trim(), expected_output); + assert_eq!(output, expected_output); }