New package prioritization and class PrevalenceCounter

1 files changed, 20 insertions, 0 deletions

diff --git a/tests/test_prioritization.py b/tests/test_prioritization.py
new file mode 100644
index 00000000..3046456e
--- /dev/null
+++ b/tests/test_prioritization.py
@@ -0,0 +1,20 @@
+from pgcli.packages.prioritization import PrevalenceCounter
+
+
+def test_prevalence_counter():
+    counter = PrevalenceCounter()
+    sql = '''SELECT * FROM foo WHERE bar GROUP BY baz;
+             select * from foo;
+             SELECT * FROM foo WHERE bar GROUP
+             BY baz'''
+    counter.update(sql)
+
+    keywords = ['SELECT', 'FROM', 'GROUP BY']
+    expected = [3, 3, 2]
+    kw_counts = [counter.keyword_count(x) for x in keywords]
+    assert kw_counts == expected
+    assert counter.keyword_count('NOSUCHKEYWORD') == 0
+
+    names = ['foo', 'bar', 'baz']
+    name_counts = [counter.name_count(x) for x in names]
+    assert name_counts == [3, 2, 2]