mirror of https://github.com/dbcli/pgcli
80 lines
1.9 KiB
Python
80 lines
1.9 KiB
Python
import pytest
|
|
from unittest.mock import Mock
|
|
|
|
from pgcli.main import PGCli
|
|
|
|
|
|
# We need this fixtures because we need PGCli object to be created
|
|
# after test collection so it has config loaded from temp directory
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def default_pgcli_obj():
|
|
return PGCli()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def DEFAULT(default_pgcli_obj):
|
|
return default_pgcli_obj.row_limit
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def LIMIT(DEFAULT):
|
|
return DEFAULT + 1000
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def over_default(DEFAULT):
|
|
over_default_cursor = Mock()
|
|
over_default_cursor.configure_mock(rowcount=DEFAULT + 10)
|
|
return over_default_cursor
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def over_limit(LIMIT):
|
|
over_limit_cursor = Mock()
|
|
over_limit_cursor.configure_mock(rowcount=LIMIT + 10)
|
|
return over_limit_cursor
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def low_count():
|
|
low_count_cursor = Mock()
|
|
low_count_cursor.configure_mock(rowcount=1)
|
|
return low_count_cursor
|
|
|
|
|
|
def test_row_limit_with_LIMIT_clause(LIMIT, over_limit):
|
|
cli = PGCli(row_limit=LIMIT)
|
|
stmt = "SELECT * FROM students LIMIT 1000"
|
|
|
|
result = cli._should_limit_output(stmt, over_limit)
|
|
assert result is False
|
|
|
|
cli = PGCli(row_limit=0)
|
|
result = cli._should_limit_output(stmt, over_limit)
|
|
assert result is False
|
|
|
|
|
|
def test_row_limit_without_LIMIT_clause(LIMIT, over_limit):
|
|
cli = PGCli(row_limit=LIMIT)
|
|
stmt = "SELECT * FROM students"
|
|
|
|
result = cli._should_limit_output(stmt, over_limit)
|
|
assert result is True
|
|
|
|
cli = PGCli(row_limit=0)
|
|
result = cli._should_limit_output(stmt, over_limit)
|
|
assert result is False
|
|
|
|
|
|
def test_row_limit_on_non_select(over_limit):
|
|
cli = PGCli()
|
|
stmt = "UPDATE students SET name='Boby'"
|
|
result = cli._should_limit_output(stmt, over_limit)
|
|
assert result is False
|
|
|
|
cli = PGCli(row_limit=0)
|
|
result = cli._should_limit_output(stmt, over_limit)
|
|
assert result is False
|