Code refactoring, test case for QueryThread added

AvivoreXT/Avivore.py

@@ -1,6 +1,6 @@
 # -*- coding: utf8 -*-
 
-from AvivoreXT import AvivoreConfig, Compat, Helper
+from AvivoreXT import AvivoreConfig, Compat, Helper, AvivoreTwitterError
 from twitter import *
 from twitter.stream import Timeout, HeartbeatTimeout, Hangup
 if Compat.is_python3():
@@ -31,9 +31,9 @@ def twitter_auth(self):
                                                      self.avivore_config.twitter_consumer_secret)
             self.twitter_instance = Twitter(auth=OAuth2(bearer_token=self.twitter_bearer_token))
         except (TwitterHTTPError, URLError):
-            SystemExit(1, "[!] Can't authenticate, check your network connection!")
+            raise AvivoreTwitterError.TwitterAuthenticationException('Network error, check your connection!')
         except Exception as e:
-            SystemExit(1, "[!] Unknown error:\n"+e)
+            raise AvivoreTwitterError.TwitterAuthenticationException(e)
 
         return self.twitter_instance
 
@@ -103,7 +103,7 @@ def twitter_search(self, search_string):
             Helper.output("[!] Problem connecting to twitter...")
             output = None  # If this bombs out, we have the option of at least spitting out a result.
         except Exception as e:
-            Helper.output("[!] Unknown problem querying twitter:\n"+e)
+            Helper.output("[!] Unknown problem querying twitter:\n%s" % e)
             output = None  # If this bombs out, we have the option of at least spitting out a result.
         return output
 
@@ -168,12 +168,24 @@ def __db_dup_check(self, tid, value):
         return output
 
     def __db_write_value(self, stime, stype, user, userid, value, tweetid, message):
+        """
+        Writes a data set found in a tweet to the 'Data' table in the result database.
+
+        :param stime:   Unix timestamp of the time the tweet was found
+        :param stype:   Type identifier mapping the tweet to the type definitions in the configuration
+        :param user:    User name of the creator of the tweet
+        :param userid:  User id of the creator of the tweet
+        :param value:   Extracted information (according to type definition)
+        :param tweetid: Tweet id
+        :param message: The tweets message itself
+        :return:        None
+        """
         # Just a simple function to write the results to the database.
         con = lite.connect(self.avivore_config.database_path)
         qstring = "INSERT INTO Data VALUES(?, ?, ?, ?, ?, ?, ?)"
         with con:
             cur = con.cursor()
-            cur.execute(qstring,
-                        (unicode(stime), unicode(stype), unicode(user), unicode(userid),
-                         unicode(value), unicode(tweetid), unicode(message)))
+            cur.execute(qstring, (stime, stype, user, userid, value, tweetid, message))
+                        #(unicode(stime), unicode(stype), unicode(user), unicode(userid),
+                         #unicode(value), unicode(tweetid), unicode(message)))
             # lid = cur.lastrowid

AvivoreXT/AvivoreMain.py

@@ -11,6 +11,12 @@
 
 
 def main(argv):
+    """
+    Avivore main() function.
+
+    :param argv:    Array containing command line arguments.
+    :return:        None
+    """
     # create AvivoreConfig instance
     if "-c" == argv[1]:
         config = AvivoreConfig.AvivoreConfig(0, argv[2])
@@ -41,16 +47,35 @@ def main(argv):
 
 
 def software_init_msg(version):
+    """
+    Prints AvivoreXT welcome message.
+
+    :param version: Version number of AvivoreXT release.
+    :return:        None
+    """
     print("AvivoreXT " + version + " by rc0r (https://github.com/rc0r)")
 
 
 def software_exit(status, message):
+    """
+    Exit function, prints exit message and quits AvivoreXT with a user provided status code.
+
+    :param status:  Exit status code.
+    :param message: Message to display.
+    :return:        None
+    """
     print("")
     print(message)
     sys.exit(status)
 
 
 def check_usage(argv):
+    """
+    Checks for correct program invocation from the command line and prints usage information if necessary.
+
+    :param argv:    Array containing all command line arguments passed to AvivoreXT
+    :return:        None
+    """
     print_usage = False
     if 3 != len(argv):
         print_usage = True
@@ -64,6 +89,11 @@ def check_usage(argv):
 
 
 def start():
+    """
+    AvivoreXT start procedure.
+
+    :return:    None
+    """
     software_init_msg(AvivoreXT.__version__)
     check_usage(sys.argv)
     try:

AvivoreXT/AvivoreTwitterError.py

@@ -0,0 +1,20 @@
+# -*- coding: utf8 -*-
+
+
+class TwitterException(Exception):
+    """
+    Generic Twitter exception class
+    """
+    def __init__(self, msg):
+        self.msg = msg
+        Exception.__init__(self, '%s' % msg)
+
+
+class TwitterAuthenticationException(TwitterException):
+    """
+    Twitter exception class for authentication failures.
+    """
+    def __init__(self, reason):
+        self.msg = 'Twitter authentication failure! Reason: %s' % (reason)
+        TwitterException.__init__(self, self.msg)
+

AvivoreXT/QueryThread.py

@@ -1,36 +1,68 @@
 # -*- coding: utf8 -*-
-from AvivoreXT import Avivore, Helper
+from AvivoreXT import Helper
 import threading
 import time
 
 
 class QueryThread(threading.Thread):
+    """
+    Class implementing the thread that performs periodic twitter searches.
+    """
     def __init__(self, avivore):
         self.avivore = avivore
+        self.stored = []
         threading.Thread.__init__(self)
 
-    def run(self):
+    def init_run(self):
+        """
+        Prints info message and performs application authentication with Twitter.
+
+        :return:        None
+        """
         Helper.output("Spawning query thread [Q].")
-        self.twitter_query_main(self.avivore)
+        self.avivore.twitter_auth()
+
+    def run(self):
+        """
+        Contains threads main loop that runs until the user kills the AvivoreXT process or some serious error occurs.
 
-    # performs classic search queries with configured time interval
-    def twitter_query_main(self, avivore):
-        stored = []
-        avivore.twitter_auth()
+        :return:        None
+        """
+        self.init_run()
         while 1:
-            for x in avivore.avivore_config.twitter_search_terms:
+            self.twitter_search(self.avivore)
+
+    def extract_data_from_tweet(self, avivore, tweet):
+        """
+        Scans data for extractable data sets (previously defined in the 'TypeDefs' section of the AvivoreXT
+        configuration) and if found stores them in the result database.
+
+        :param avivore: Previously initialized Avivore instance.
+        :param tweet:   The tweet to analyzed that was returned from Twitter as a JSON object.
+        :return:        None
+        """
+        z = tweet['id'], tweet['created_at'], tweet['user']['screen_name'], tweet['text'], tweet['user']['id_str']
+        result = avivore.twitter_read_tweet(z[3])
+        if result[0] >= 0:  # If something is found, then we'll process the tweet
+            self.stored = self.stored, int(z[0])
+            # result value, time, result itself, tweet ID, tweet itself, userId
+            string = result[0], z[2], result[1], z[0], z[3], z[4]
+            message = avivore.process_tweet(string)
+            if 0 != message:
+                Helper.output("[Q] " + message)
+
+    def twitter_search(self, avivore):
+        """
+        Performs a Twitter search query for all keywords defined in 'csv_search_term' configuration option and returns
+        all found tweets.
+
+        :param avivore: Previously initialized Avivore instance.
+        :return:        None
+        """
+        # performs classic search queries with configured time interval
+        for x in avivore.avivore_config.twitter_search_terms:
                 twit_data = avivore.twitter_search(x)
                 if twit_data is not None:
                     for y in avivore.twitter_search(x):
-                        z = y['id'], y['created_at'], y['user']['screen_name'], y['text'], y['user']['id_str']
-                        result = avivore.twitter_read_tweet(z[3])
-                        if result[0] < 0:
-                            pass
-                        else:  # If something is found, then we'll process the tweet
-                            stored = stored, int(z[0])
-                            # result value, time, result itself, tweet ID, tweet itself, userId
-                            string = result[0], z[2], result[1], z[0], z[3], z[4]
-                            message = avivore.process_tweet(string)
-                            if 0 != message:
-                                Helper.output("[Q] " + message)
+                        self.extract_data_from_tweet(avivore, y)
                 time.sleep(avivore.avivore_config.twitter_search_interval)

test/test_QueryThread.py

@@ -1 +1,59 @@
-__author__ = 'rc0r'
+# -*- coding: utf-8 -*-
+
+import unittest
+from AvivoreXT import Avivore, QueryThread, AvivoreTwitterError
+import os
+
+
+class QueryThreadTestCase(unittest.TestCase):
+    def setUp(self):
+        # code to execute in preparation for tests
+        os.symlink('../sample_conf.db', './testdata/test.db')
+
+        self.avivore_conf = Avivore.AvivoreConfig.AvivoreConfig(1, './testdata/test.db')
+        self.avivore_conf.read_config()
+        self.avivore_conf.init_database()
+        self.avivore_conf.twitter_search_interval = 0
+        self.avivore = Avivore.Avivore(self.avivore_conf)
+        self.query_thread_instance = QueryThread.QueryThread(self.avivore)
+
+    def tearDown(self):
+        # code to execute to clean up after tests
+        os.remove('./testdata/test.db')
+
+    def test_run_init(self):
+        # test for network failure or twitter authentication error (no valid twitter credentials in config file)
+        with self.assertRaises(AvivoreTwitterError.TwitterAuthenticationException):
+            self.assertIsNone(self.query_thread_instance.init_run())
+
+    @unittest.skip('This test assumes there are valid twitter credentials in sample_conf.db!')
+    def test_run_init_working_creds(self):
+        # test in case network is working
+        self.assertIsNone(self.query_thread_instance.init_run())
+
+    def test_twitter_search(self):
+        self.assertIsNone(self.query_thread_instance.twitter_search(self.avivore))
+
+    def test_extract_data_from_tweet(self):
+        tweet = {
+            'id': '0123456789',
+            'created_at': '000000000',
+            'user': {
+                'id_str': '987654321',
+                'screen_name': '_rc0r',
+            },
+            'text': 'This is a test tweet! My server ip: 10.11.12.13!',
+        }
+
+        # process unseen tweet
+        self.assertIsNone(self.query_thread_instance.extract_data_from_tweet(self.avivore, tweet))
+        # process dupe
+        self.assertIsNone(self.query_thread_instance.extract_data_from_tweet(self.avivore, tweet))
+
+
+def test_main():
+    unittest.main()
+
+
+if __name__ == '__main__':
+    test_main()