def run_parsing(config): parser = argparse.ArgumentParser() subparsers = parser.add_subparsers() # Search parser group search_parser = subparsers.add_parser('search') search_parser.add_argument('QUERY', action='store', type=str, default=False) search_parser.add_argument('--output', '-o', action='store', required=False, default='json', choices=['tsv', 'json', 'binary']) search_parser.add_argument('--columns', '-c', action='store', required=False, help='Coma separated names of columns') # Prune parser group prune_parser = subparsers.add_parser('prune') prune_parser.add_argument('PRUNE', action='store', type=int, default=False) args = parser.parse_args() try: db_path = os.environ['DB_PATH'] except KeyError: print("Env variable DB_PATH is missing!", file=sys.stderr) exit(1) db = Database(db_path) if 'QUERY' in args: try: # Scrape and download down = Downloader() parser = Parser(config) d = Driver(config, args.QUERY, down, parser, db) torrents, headers = d.run() # Format and print fm = lib.utils.Formatter(headers, torrents) if args.columns: fm.filter_data(args.columns.split(',')) fm.format_data(args.output) fm.print_data() except: traceback.print_exc() exit(1) elif 'PRUNE' in args: try: # Prune DB db.connect() db.prune(args.PRUNE, time.time()) except: traceback.print_exc() exit(1)
class TestDBPrune(unittest.TestCase): def setUp(self) -> None: self.db = Database(':memory:') self.db.full_db_path = ':memory:' self.db.connect() populate_db(self.db.cur) def tearDown(self) -> None: self.db.close_conn() def test_delete_right_query(self): self.db.prune(200, NOW) self.assertEqual([MOCK_DATA['queries'][0]], self.db.cur.execute('SELECT * FROM queries') .fetchall()) def test_delete_correct_binding(self): self.db.prune(200, NOW) self.assertEqual([MOCK_DATA['query2torrent'][0], MOCK_DATA['query2torrent'][1], MOCK_DATA['query2torrent'][2]], self.db.cur.execute('SELECT * FROM query2torrent') .fetchall()) def test_delete_correct_torrents(self): self.db.prune(200, NOW) self.assertEqual([MOCK_DATA['torrents'][0], MOCK_DATA['torrents'][1], MOCK_DATA['torrents'][4]], self.db.cur.execute('SELECT id, name FROM torrents') .fetchall()) def test_do_not_delete(self): self.db.prune(301, NOW) self.assertEqual(MOCK_DATA['queries'], self.db.cur.execute('SELECT * FROM queries') .fetchall())