integrational_tests.py

"""
This test module will execute real commands using shell.
This means it will call sgpt.py with command line arguments.
Make sure you have your API key in place ~/.config/shell-gpt/api_key.txt
It is useful for quick tests, saves a bit time.
"""

import subprocess
import os
from time import sleep
from unittest import TestCase
from tempfile import NamedTemporaryFile

import typer
from typer.testing import CliRunner
from sgpt import main

runner = CliRunner()
app = typer.Typer()
app.command()(main)


class TestCliApp(TestCase):
    def setUp(self) -> None:
        # Just to not spam the API.
        sleep(2)

    @staticmethod
    def get_arguments(prompt, **kwargs):
        arguments = [prompt]
        for key, value in kwargs.items():
            arguments.append(key)
            if isinstance(value, bool):
                continue
            arguments.append(value)
        # arguments.extend([kv for kv in kwargs.items() for kv in kv])
        return arguments

    def test_simple_queries(self):
        dict_arguments = {
            "prompt": "What is the capital of the Czech Republic?",
            "--max-tokens": 32,
            "--model": "curie",
        }
        result = runner.invoke(app, self.get_arguments(**dict_arguments))
        assert result.exit_code == 0
        assert "Prague" in result.stdout

    def test_shell_queries(self):
        dict_arguments = {
            "prompt": "make a commit using git",
            "--max-tokens": 32,
            "--model": "davinci",
            "--shell": True,
        }
        result = runner.invoke(app, self.get_arguments(**dict_arguments))
        assert result.exit_code == 0
        assert "git commit" in result.stdout

    def test_code_queries(self):
        """
        This test will request from GPT-3 a python code to make CLI app,
        which will be written to a temp file, and then it will be executed
        in shell with two positional int arguments. As the output we are
        expecting the result of multiplying them.
        """
        dict_arguments = {
            "prompt": (
                "Create a command line application using Python that "
                "accepts two integer positional command line arguments "
                "and prints the result of multiplying them."
            ),
            "--max-tokens": 128,
            "--model": "davinci",
            "--code": True,
        }
        result = runner.invoke(app, self.get_arguments(**dict_arguments))
        assert result.exit_code == 0
        # Since output will be slightly different, there is no way how to test it precisely.
        assert "print" in result.stdout
        assert "*" in result.stdout
        with NamedTemporaryFile("w+", delete=False) as file:
            try:
                compile(result.output, file.name, "exec")
            except SyntaxError:
                assert False, "The output is not valid Python code."
            file.seek(0)
            file.truncate()
            file.write(result.output)
            file_path = file.name
        number_a = number_b = 2
        # Execute output code in the shell with arguments.
        arguments = ["python", file.name, str(number_a), str(number_b)]
        script_output = subprocess.run(arguments, stdout=subprocess.PIPE, check=True)
        os.remove(file_path)
        assert script_output.stdout.decode().strip(), number_a * number_b