pgcli/tests/test_init_commands_simple.py at a98163ecf82d920735a8289ec9269a53dbbf30fe · dbcli/pgcli · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import pytest
from click.testing import CliRunner

from pgcli.main import cli, PGCli


@pytest.fixture
def dummy_exec(monkeypatch, tmp_path):
    # Capture executed commands
    # Isolate config directory for tests
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path))
    dummy_cmds = []

    class DummyExec:
        def run(self, cmd):
            # Ignore ping SELECT 1 commands used for exiting CLI
            if cmd.strip().upper() == "SELECT 1":
                return []
            # Record init commands
            dummy_cmds.append(cmd)
            return []

        def get_timezone(self):
            return "UTC"

        def set_timezone(self, *args, **kwargs):
            pass

    def fake_connect(self, *args, **kwargs):
        self.pgexecute = DummyExec()

    monkeypatch.setattr(PGCli, "connect", fake_connect)
    return dummy_cmds


def test_init_command_option(dummy_exec):
    "Test that --init-command triggers execution of the command."
    runner = CliRunner()
    # Use a custom init command and --ping to exit the CLI after init commands
    result = runner.invoke(
        cli, ["--init-command", "SELECT foo", "--ping", "db", "user"]
    )
    assert result.exit_code == 0
    # Should print the init command
    assert "Running init commands: SELECT foo" in result.output
    # Should exit via ping
    assert "PONG" in result.output
    # DummyExec should have recorded only the init command
    assert dummy_exec == ["SELECT foo"]


def test_init_commands_from_config(dummy_exec, tmp_path):
    """
    Test that init commands defined in the config file are executed on startup.
    """
    # Create a temporary config file with init-commands
    config_file = tmp_path / "pgclirc_test"
    config_file.write_text(
        "[main]\n[init-commands]\nfirst = SELECT foo;\nsecond = SELECT bar;\n"
    )

    runner = CliRunner()
    # Use --ping to exit the CLI after init commands
    result = runner.invoke(
        cli, ["--pgclirc", str(config_file.absolute()), "--ping", "testdb", "user"]
    )
    assert result.exit_code == 0
    # Should print both init commands in order (note trailing semicolons cause double ';;')
    assert "Running init commands: SELECT foo;; SELECT bar;" in result.output
    # DummyExec should have recorded both commands
    assert dummy_exec == ["SELECT foo;", "SELECT bar;"]


def test_init_commands_option_and_config(dummy_exec, tmp_path):
    """
    Test that CLI-provided init command is appended after config-defined commands.
    """
    # Create a temporary config file with init-commands
    config_file = tmp_path / "pgclirc_test"
    config_file.write_text("[main]\n [init-commands]\nfirst = SELECT foo;\n")

    runner = CliRunner()
    # Use --ping to exit the CLI after init commands
    result = runner.invoke(
        cli,
        [
            "--pgclirc",
            str(config_file),
            "--init-command",
            "SELECT baz;",
            "--ping",
            "testdb",
            "user",
        ],
    )
    assert result.exit_code == 0
    # Should print config command followed by CLI option (double ';' between commands)
    assert "Running init commands: SELECT foo;; SELECT baz;" in result.output
    # DummyExec should record both commands in order
    assert dummy_exec == ["SELECT foo;", "SELECT baz;"]