1 files changed, 18 insertions, 21 deletions
diff --git a/tests/test_rowlimit.py b/tests/test_rowlimit.py
index bc66f1e7..e76ea049 100644
--- a/tests/test_rowlimit.py
+++ b/tests/test_rowlimit.py
@@ -1,6 +1,7 @@
-from pgcli.main import PGCli
-from mock import Mock
 import pytest
+from mock import Mock
+
+from pgcli.main import PGCli
 
 
 # We need this fixtures beacause we need PGCli object to be created
@@ -43,40 +44,36 @@ def low_count():
     return low_count_cursor
 
 
-def test_default_row_limit(low_count, over_default):
-    cli = PGCli()
-    stmt = "SELECT * FROM students"
-    result = cli._should_show_limit_prompt(stmt, low_count)
+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
 
-    result = cli._should_show_limit_prompt(stmt, over_default)
-    assert result is True
+    cli = PGCli(row_limit=0)
+    result = cli._should_limit_output(stmt, over_limit)
+    assert result is False
 
 
-def test_set_row_limit(over_default, over_limit, LIMIT):
+def test_row_limit_without_LIMIT_clause(LIMIT, over_limit):
     cli = PGCli(row_limit=LIMIT)
     stmt = "SELECT * FROM students"
-    result = cli._should_show_limit_prompt(stmt, over_default)
-    assert result is False
 
-    result = cli._should_show_limit_prompt(stmt, over_limit)
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is True
 
-
-def test_no_limit(over_limit):
     cli = PGCli(row_limit=0)
-    stmt = "SELECT * FROM students"
-
-    result = cli._should_show_limit_prompt(stmt, over_limit)
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is False
 
 
-def test_row_limit_on_non_select(over_default):
+def test_row_limit_on_non_select(over_limit):
     cli = PGCli()
-    stmt = "UPDATE students set name='Boby'"
-    result = cli._should_show_limit_prompt(stmt, None)
+    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_show_limit_prompt(stmt, over_default)
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is False