From 5ec5b60660d33bda30d633a7089f1070c546d768 Mon Sep 17 00:00:00 2001 From: Paul Gauthier Date: Thu, 6 Jul 2023 11:50:13 -0700 Subject: [PATCH] better /tokens test --- tests/test_commands.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_commands.py b/tests/test_commands.py index 88e292eb4..7645fef13 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -1,12 +1,15 @@ import codecs import os import shutil +import sys import tempfile +from io import StringIO from unittest import TestCase -import io from aider import models +from aider.coders import Coder from aider.commands import Commands +from aider.dump import dump # noqa: F401 from aider.io import InputOutput @@ -54,7 +57,6 @@ class TestCommands(TestCase): def test_cmd_tokens(self): # Initialize the Commands and InputOutput objects io = InputOutput(pretty=False, yes=True) - from aider.coders import Coder coder = Coder.create(models.GPT35, None, io, openai_api_key="deadbeef") commands = Commands(io, coder) @@ -62,7 +64,7 @@ class TestCommands(TestCase): commands.cmd_add("foo.txt bar.txt") # Redirect the standard output to an instance of io.StringIO - stdout = io.StringIO() + stdout = StringIO() sys.stdout = stdout commands.cmd_tokens("") @@ -72,4 +74,6 @@ class TestCommands(TestCase): # Get the console output console_output = stdout.getvalue() - print(console_output) + + self.assertIn("foo.txt", console_output) + self.assertIn("bar.txt", console_output)