From 30406483b9a37049c493eafb1b3554f64c217c2a Mon Sep 17 00:00:00 2001
From: Ben Kurtovic <ben.kurtovic@gmail.com>
Date: Wed, 14 May 2014 13:16:58 -0400
Subject: [PATCH] Add unit tests, cleanup.

---
 bitshift/languages.py       |  1 -
 bitshift/parser/__init__.py |  2 +-
 test/test_query_parser.py   | 37 +++++++++++++++++++++++++++++++++----
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/bitshift/languages.py b/bitshift/languages.py
index b04c094..6b89c12 100644
--- a/bitshift/languages.py
+++ b/bitshift/languages.py
@@ -1,2 +1 @@
-
 LANGS = ["Python", "C", "Java", "Ruby"]
diff --git a/bitshift/parser/__init__.py b/bitshift/parser/__init__.py
index 55c76e1..bc22514 100644
--- a/bitshift/parser/__init__.py
+++ b/bitshift/parser/__init__.py
@@ -21,7 +21,7 @@ def _lang(codelet):
 
     if codelet.filename is not None:
         try:
-            return pgl.guess_lexer_for_filename(codelet.filename, '').name
+            return pgl.guess_lexer_for_filename(codelet.filename, codelet.code).name
         except:
             raise UnsupportedFileError('Could not find a lexer for the codelet\'s filename')
 
diff --git a/test/test_query_parser.py b/test/test_query_parser.py
index c74da4a..6b33813 100644
--- a/test/test_query_parser.py
+++ b/test/test_query_parser.py
@@ -3,14 +3,43 @@ import unittest
 
 from bitshift.query import parse_query
 
+TESTS = [
+    # Text
+    ("test", "Tree(Text(String(u'test')))"),
+    ("re:test", "Tree(Text(Regex(u'test')))"),
+
+    # Language
+    ("language:python", "Tree(Language(Python))"),
+    ("language:py", "Tree(Language(Python))"),
+    ("l:r:r..y", "Tree(Language(Ruby))"),
+    ("lang:re:py|c", "Tree(BinaryOp(Language(C), OR, Language(Python)))"),
+
+    # Author
+    ('"author:Ben Kurtovic"', "Tree(Author(String(u'Ben Kurtovic')))"),
+    (r"'a:re:b.*?\sk.*?'", r"Tree(Author(Regex(u'b.*?\\sk.*?')))"),
+
+    # Date
+    ("'create:before:Jan 1, 2014'",
+     "Tree(Date(CREATE, BEFORE, 2014-01-01 00:00:00))"),
+    ("'modify:after:2010-05-09 10:11:12'",
+     "Tree(Date(MODIFY, AFTER, 2010-05-09 10:11:12))"),
+
+    # Symbol
+    ("sym:foobar", "Tree(Symbol(ALL, String(u'foobar')))"),
+    ("func:foo_bar", "Tree(Symbol(FUNCTION, String(u'foo_bar')))"),
+    ("func:foo_bar()", "Tree(Symbol(FUNCTION, String(u'foo_bar')))"),
+    ("class:FooBar", "Tree(Symbol(CLASS, String(u'FooBar')))"),
+    ("var:foobar", "Tree(Symbol(VARIABLE, String(u'foobar')))"),
+    ("var:r:foobar", "Tree(Symbol(VARIABLE, Regex(u'foobar')))"),
+]
+
 class TestQueryParser(unittest.TestCase):
     """Unit tests for the query parser in :py:mod:`bitshift.query`."""
 
     def test_parse(self):
-        """test basic query parsing"""
-        pq = lambda s: parse_query(s).serialize()
-        self.assertEqual("Tree(Text(String('test')))", pq("test"))
-        self.assertEqual("Tree(Text(Regex('test')))", pq("re:test"))
+        """test full query parsing"""
+        for test, expected in TESTS:
+            self.assertEqual(expected, parse_query(test).serialize())
 
 
 if __name__ == "__main__":