diff --git a/Tests/GoogleAITests/GenerateContentResponses/unary-failure-unknown-enum-finish-reason.json b/Tests/GoogleAITests/GenerateContentResponses/unary-failure-unknown-enum-finish-reason.json new file mode 100644 index 0000000..1adbc70 --- /dev/null +++ b/Tests/GoogleAITests/GenerateContentResponses/unary-failure-unknown-enum-finish-reason.json @@ -0,0 +1,53 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Some text" + } + ] + }, + "finishReason": "FAKE_NEW_FINISH_REASON", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/Tests/GoogleAITests/GenerateContentResponses/unary-failure-unknown-enum-prompt-blocked.json b/Tests/GoogleAITests/GenerateContentResponses/unary-failure-unknown-enum-prompt-blocked.json new file mode 100644 index 0000000..a00cbb7 --- /dev/null +++ b/Tests/GoogleAITests/GenerateContentResponses/unary-failure-unknown-enum-prompt-blocked.json @@ -0,0 +1,23 @@ +{ + "promptFeedback": { + "blockReason": "FAKE_NEW_BLOCK_REASON", + "safetyRatings": [ + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HATE_SPEECH", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE" + } + ] + } +} diff --git a/Tests/GoogleAITests/GenerateContentResponses/unary-success-unknown-enum-safety-ratings.json b/Tests/GoogleAITests/GenerateContentResponses/unary-success-unknown-enum-safety-ratings.json new file mode 100644 index 0000000..3c47a21 --- /dev/null +++ b/Tests/GoogleAITests/GenerateContentResponses/unary-success-unknown-enum-safety-ratings.json @@ -0,0 +1,46 @@ +{ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "Some text" + } + ], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [ + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "MEDIUM" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "FAKE_NEW_HARM_PROBABILITY" + }, + { + "category": "FAKE_NEW_HARM_CATEGORY", + "probability": "HIGH" + } + ] + } + ], + "promptFeedback": { + "safetyRatings": [ + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "MEDIUM" + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "FAKE_NEW_HARM_PROBABILITY" + }, + { + "category": "FAKE_NEW_HARM_CATEGORY", + "probability": "HIGH" + } + ] + } +} diff --git a/Tests/GoogleAITests/GenerateContentResponses/unary-success-unknown-enum.json b/Tests/GoogleAITests/GenerateContentResponses/unary-success-unknown-enum.json deleted file mode 100644 index b27a11a..0000000 --- a/Tests/GoogleAITests/GenerateContentResponses/unary-success-unknown-enum.json +++ /dev/null @@ -1,52 +0,0 @@ -{ - "candidates": [ - { - "content": { - "parts": [ - { - "text": "1. **Use Freshly Ground Coffee**:\n - Grind your coffee beans just before brewing to preserve their flavor and aroma.\n - Use a burr grinder for a consistent grind size.\n\n\n2. **Choose the Right Water**:\n - Use filtered or spring water for the best taste.\n - Avoid using tap water, as it may contain impurities that can affect the flavor.\n\n\n3. **Measure Accurately**:\n - Use a kitchen scale to measure your coffee and water precisely.\n - A general rule of thumb is to use 1:16 ratio of coffee to water (e.g., 15 grams of coffee to 240 grams of water).\n\n\n4. **Preheat Your Equipment**:\n - Preheat your coffee maker or espresso machine before brewing to ensure a consistent temperature.\n\n\n5. **Control the Water Temperature**:\n - The ideal water temperature for brewing coffee is between 195°F (90°C) and 205°F (96°C).\n - Too hot water can extract bitter flavors, while too cold water won't extract enough flavor.\n\n\n6. **Steep the Coffee**:\n - For drip coffee, let the water slowly drip through the coffee grounds for optimal extraction.\n - For pour-over coffee, pour the water in a circular motion over the coffee grounds, allowing it to steep for 30-45 seconds before continuing.\n\n\n7. **Clean Your Equipment**:\n - Regularly clean your coffee maker or espresso machine to prevent the buildup of oils and residue that can affect the taste of your coffee.\n\n\n8. **Experiment with Different Coffee Beans**:\n - Try different coffee beans from various regions and roasts to find your preferred flavor profile.\n - Experiment with different grind sizes and brewing methods to optimize the flavor of your chosen beans.\n\n\n9. **Store Coffee Properly**:\n - Store your coffee beans in an airtight container in a cool, dark place to preserve their freshness and flavor.\n - Avoid storing coffee in the refrigerator or freezer, as this can cause condensation and affect the taste.\n\n\n10. **Enjoy Freshly Brewed Coffee**:\n - Drink your coffee as soon as possible after brewing to enjoy its peak flavor and aroma.\n - Coffee starts to lose its flavor and aroma within 30 minutes of brewing." - } - ] - }, - "index": 0, - "safetyRatings": [ - { - "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT_ENUM_NEW", - "probability": "NEGLIGIBLE" - }, - { - "category": "HARM_CATEGORY_HATE_SPEECH", - "probability": "NEGLIGIBLE" - }, - { - "category": "HARM_CATEGORY_HARASSMENT", - "probability": "NEGLIGIBLE" - }, - { - "category": "HARM_CATEGORY_DANGEROUS_CONTENT", - "probability": "NEGLIGIBLE" - } - ] - } - ], - "promptFeedback": { - "safetyRatings": [ - { - "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", - "probability": "NEGLIGIBLE" - }, - { - "category": "HARM_CATEGORY_HATE_SPEECH", - "probability": "NEGLIGIBLE" - }, - { - "category": "HARM_CATEGORY_HARASSMENT", - "probability": "NEGLIGIBLE" - }, - { - "category": "HARM_CATEGORY_DANGEROUS_CONTENT_LIKE_A_NEW_ENUM", - "probability": "NEGLIGIBLE_NEW_ENUM" - } - ] - } -} diff --git a/Tests/GoogleAITests/GenerativeModelTests.swift b/Tests/GoogleAITests/GenerativeModelTests.swift index 85bf9b2..95b0f85 100644 --- a/Tests/GoogleAITests/GenerativeModelTests.swift +++ b/Tests/GoogleAITests/GenerativeModelTests.swift @@ -134,17 +134,23 @@ final class GenerativeModelTests: XCTestCase { XCTAssertEqual(promptFeedback.safetyRatings, safetyRatingsNegligible) } - func testGenerateContent_success_unknownEnum() async throws { + func testGenerateContent_success_unknownEnum_safetyRatings() async throws { + let expectedSafetyRatings = [ + SafetyRating(category: .harassment, probability: .medium), + SafetyRating(category: .dangerousContent, probability: .unknown), + SafetyRating(category: .unknown, probability: .high), + ] MockURLProtocol .requestHandler = try httpRequestHandler( - forResource: "unary-success-unknown-enum", + forResource: "unary-success-unknown-enum-safety-ratings", withExtension: "json" ) - let content = try await model.generateContent(testPrompt) + let response = try await model.generateContent(testPrompt) - XCTAssertNotNil(content.text) - // TODO: Add assertions + XCTAssertEqual(response.text, "Some text") + XCTAssertEqual(response.candidates.first?.safetyRatings, expectedSafetyRatings) + XCTAssertEqual(response.promptFeedback?.safetyRatings, expectedSafetyRatings) } func testGenerateContent_failure_invalidAPIKey() async throws { @@ -261,7 +267,43 @@ final class GenerativeModelTests: XCTestCase { } catch let GenerateContentError.promptBlocked(response) { XCTAssertNil(response.text) } catch { - XCTFail("Should throw a promptBlocked]") + XCTFail("Should throw a promptBlocked") + } + } + + func testGenerateContent_failure_unknownEnum_finishReason() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unknown-enum-finish-reason", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.responseStoppedEarly(reason, response) { + XCTAssertEqual(reason, .unknown) + XCTAssertEqual(response.text, "Some text") + } catch { + XCTFail("Should throw a responseStoppedEarly") + } + } + + func testGenerateContent_failure_unknownEnum_promptBlocked() async throws { + MockURLProtocol + .requestHandler = try httpRequestHandler( + forResource: "unary-failure-unknown-enum-prompt-blocked", + withExtension: "json" + ) + + do { + _ = try await model.generateContent(testPrompt) + XCTFail("Should throw") + } catch let GenerateContentError.promptBlocked(response) { + let promptFeedback = try XCTUnwrap(response.promptFeedback) + XCTAssertEqual(promptFeedback.blockReason, .unknown) + } catch { + XCTFail("Should throw a promptBlocked") } }