update tests

Author: sjrl

test/token_counters/test_tiktoken_counter.py

@@ -109,6 +109,28 @@ def test_serde_round_trip(self):
         assert restored.tokens_per_image == 200
 
 
+@tool
+def search(query: Annotated[str, "the search query"]) -> str:
+    """Search the web for a query and return the top results."""
+    return "result"
+
+
+class TestTiktokenCounterTools:
+    def test_tool_schemas_add_to_the_count(self, fake_encoder):
+        # A provider is sent the schemas alongside the messages, so they consume tokens too.
+        counter = TiktokenCounter()
+        messages = [ChatMessage.from_user("hi")]
+
+        assert counter.count(messages, tools=[search]) > counter.count(messages)
+
+    def test_tools_can_be_counted_without_messages(self, fake_encoder):
+        assert TiktokenCounter().count([], tools=[search]) > 0
+
+    def test_nothing_to_measure_is_zero(self):
+        assert TiktokenCounter().count([]) == 0
+        assert TiktokenCounter().count([], tools=None) == 0
+
+
 @pytest.mark.integration
 class TestTiktokenCounterIntegration:
     """Exercises the real encoder, which downloads its vocabulary on first use."""
@@ -137,25 +159,3 @@ def test_an_image_is_charged_at_the_flat_rate(self):
         counter = TiktokenCounter(tokens_per_image=85)
 
         assert counter.count([ChatMessage.from_user(content_parts=[IMAGE])]) > 85
-
-
-@tool
-def search(query: Annotated[str, "the search query"]) -> str:
-    """Search the web for a query and return the top results."""
-    return "result"
-
-
-class TestTiktokenCounterTools:
-    def test_tool_schemas_add_to_the_count(self):
-        # A provider is sent the schemas alongside the messages, so they consume tokens too.
-        counter = TiktokenCounter()
-        messages = [ChatMessage.from_user("hi")]
-
-        assert counter.count(messages, tools=[search]) > counter.count(messages)
-
-    def test_tools_can_be_counted_without_messages(self):
-        assert TiktokenCounter().count([], tools=[search]) > 0
-
-    def test_nothing_to_measure_is_zero(self):
-        assert TiktokenCounter().count([]) == 0
-        assert TiktokenCounter().count([], tools=None) == 0