1 files changed, 79 insertions, 0 deletions
diff --git a/tests/test_rowlimit.py b/tests/test_rowlimit.py
new file mode 100644
index 0000000..da916b4
--- /dev/null
+++ b/tests/test_rowlimit.py
@@ -0,0 +1,79 @@
+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