From a23938961f20724f1eda43cbe3b0a21ae594b8e9 Mon Sep 17 00:00:00 2001
From: Stevan Freeborn <65925598+StevanFreeborn@users.noreply.github.com>
Date: Thu, 2 Jan 2025 20:41:12 -0600
Subject: [PATCH] =?UTF-8?q?tests:=20add=20integration=20tests=20to=20cover?=
 =?UTF-8?q?=20both=20=F0=9F=98=81=20and=20=F0=9F=A5=B2=20paths=20when=20co?=
 =?UTF-8?q?unting=20message=20tokens?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../Integration/AnthropicApiClientTests.cs    | 85 +++++++++++++++++++
 .../Integration/IntegrationTest.cs            | 22 ++++-
 2 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/tests/AnthropicClient.Tests/Integration/AnthropicApiClientTests.cs b/tests/AnthropicClient.Tests/Integration/AnthropicApiClientTests.cs
index 55a22e2..f874a31 100644
--- a/tests/AnthropicClient.Tests/Integration/AnthropicApiClientTests.cs
+++ b/tests/AnthropicClient.Tests/Integration/AnthropicApiClientTests.cs
@@ -418,4 +418,89 @@ public class AnthropicApiClientTests : IntegrationTest
     textContent.As<TextContent>().Text.Should().Be("It is a PDF");
     textContent.As<TextContent>().Type.Should().Be("text");
   }
+
+  [Fact]
+  public async Task CountMessageTokensAsync_WhenCalled_ItShouldReturnCountTokensResponse()
+  {
+    _mockHttpMessageHandler
+      .WhenCountMessageTokensRequest()
+      .Respond(
+        HttpStatusCode.OK,
+        "application/json",
+        @"{
+          ""input_tokens"": 10
+        }"
+      );
+
+    var request = new CountMessageTokensRequest(
+      model: AnthropicModels.Claude35Sonnet,
+      messages: [
+        new(MessageRole.User, [new TextContent("Hello!")]),
+      ]
+    );
+
+    var result = await Client.CountMessageTokensAsync(request);
+
+    result.IsSuccess.Should().BeTrue();
+    result.Value.Should().BeOfType<TokenCountResponse>();
+    result.Value.InputTokens.Should().Be(10);
+  }
+
+  [Fact]
+  public async Task CountMessageTokensAsync_WhenCalledAndErrorReturned_ItShouldHandleError()
+  {
+    _mockHttpMessageHandler
+      .WhenCountMessageTokensRequest()
+      .Respond(
+        HttpStatusCode.BadRequest,
+        "application/json",
+        @"{
+          ""type"": ""error"",
+          ""error"": {
+            ""type"": ""invalid_request_error"",
+            ""message"": ""messages: roles must alternate between user and assistant, but found multiple user roles in a row""
+          }
+        }"
+      );
+
+    var request = new CountMessageTokensRequest(
+      model: AnthropicModels.Claude35Sonnet,
+      messages: [
+        new(MessageRole.User, [new TextContent("Hello!")]),
+        new(MessageRole.User, [new TextContent("Hello!")])
+      ]
+    );
+
+    var result = await Client.CountMessageTokensAsync(request);
+
+    result.IsSuccess.Should().BeFalse();
+    result.Error.Should().BeOfType<AnthropicError>();
+    result.Error.Error.Should().BeOfType<InvalidRequestError>();
+  }
+
+  [Fact]
+  public async Task CountMessageTokensAsync_WhenCalledRequestFailsAndCanNotSerializeError_ItShouldReturnUnknownError()
+  {
+    _mockHttpMessageHandler
+      .WhenCountMessageTokensRequest()
+      .Respond(
+        HttpStatusCode.BadRequest,
+        "application/json",
+        @"{}"
+      );
+
+    var request = new CountMessageTokensRequest(
+      model: AnthropicModels.Claude35Sonnet,
+      messages: [
+        new(MessageRole.User, [new TextContent("Hello!")]),
+        new(MessageRole.User, [new TextContent("Hello!")])
+      ]
+    );
+
+    var result = await Client.CountMessageTokensAsync(request);
+
+    result.IsSuccess.Should().BeFalse();
+    result.Error.Should().BeOfType<AnthropicError>();
+    result.Error.Error.Should().BeOfType<ApiError>();
+  }
 }
\ No newline at end of file
diff --git a/tests/AnthropicClient.Tests/Integration/IntegrationTest.cs b/tests/AnthropicClient.Tests/Integration/IntegrationTest.cs
index efc74a9..fd04cec 100644
--- a/tests/AnthropicClient.Tests/Integration/IntegrationTest.cs
+++ b/tests/AnthropicClient.Tests/Integration/IntegrationTest.cs
@@ -13,10 +13,18 @@ public class IntegrationTest
 
 public static class MockHttpMessageHandlerExtensions
 {
-  private static MockedRequest SetupBaseRequest(this MockHttpMessageHandler mockHttpMessageHandler)
+  private const string BaseUrl = "https://api.anthropic.com/v1";
+  private static readonly string MessagesEndpoint = $"{BaseUrl}/messages";
+  private static readonly string CountTokensEndpoint = $"{BaseUrl}/messages/count_tokens";
+
+  private static MockedRequest SetupBaseRequest(
+    this MockHttpMessageHandler mockHttpMessageHandler,
+    HttpMethod method,
+    string url
+  )
   {
     return mockHttpMessageHandler
-      .When(HttpMethod.Post, "https://api.anthropic.com/v1/messages")
+      .When(method, url)
       .WithHeaders(new Dictionary<string, string>
       {
         { "anthropic-version", "2023-06-01" },
@@ -27,14 +35,20 @@ public static class MockHttpMessageHandlerExtensions
   public static MockedRequest WhenCreateMessageRequest(this MockHttpMessageHandler mockHttpMessageHandler)
   {
     return mockHttpMessageHandler
-      .SetupBaseRequest()
+      .SetupBaseRequest(HttpMethod.Post, MessagesEndpoint)
       .WithJsonContent<MessageRequest>(r => r.Stream == false, JsonSerializationOptions.DefaultOptions);
   }
 
   public static MockedRequest WhenCreateStreamMessageRequest(this MockHttpMessageHandler mockHttpMessageHandler)
   {
     return mockHttpMessageHandler
-      .SetupBaseRequest()
+      .SetupBaseRequest(HttpMethod.Post, MessagesEndpoint)
       .WithJsonContent<StreamMessageRequest>(r => r.Stream == true, JsonSerializationOptions.DefaultOptions);
   }
+
+  public static MockedRequest WhenCountMessageTokensRequest(this MockHttpMessageHandler mockHttpMessageHandler)
+  {
+    return mockHttpMessageHandler
+      .SetupBaseRequest(HttpMethod.Post, CountTokensEndpoint);
+  }
 }
\ No newline at end of file