class SchedulerTest(TestCase):
@classmethod
def setUpClass(cls):
setup(['-m', 'mongodb://localhost'],
do_not_parse_config=True)
super(SchedulerTest, cls).setUpClass()
def setUp(self):
self.host = Host(host="example.com")
self.host.save()
@mock.patch("toddler.rabbit_sender.send_message_sync")
def test_scheduling(self, send_message):
from toddler.scheduler import Scheduler
s = Scheduler(rabbitmq_url="amqp://localhost", exchange='CrawlRequest',
routing_key='CrawlRequest')
s.schedule_jobs_for_hosts()
host = Host.objects(host=self.host.host).first()
self.assertEqual(
host.last_crawl_job_date.date(),
datetime.utcnow().date()
)
# self.assertTrue(send_message.called)
@mock.patch("toddler.rabbit_sender.send_message_sync")
def test_analysis_task_delay_queue(self, send_message):
from toddler.scheduler import AnalysisTaskDelayQueueObserver
a = AnalysisTaskDelayQueueObserver(rabbitmq_url="amqp://localhost",
queue='AnalysisTaskDelayQueue',
exchange="AnalysisTask",
routing_key='AnalysisTask')
d = {
"delay_reason": "test reason",
"timeout": (datetime.utcnow() - timedelta(hours=1)).isoformat(),
"message": {
"test": "test"
}
}
a.process_task(
ujson.dumps(d).encode("utf8")
)
self.assertTrue(send_message.called)
self.assertEqual(send_message.call_args[0][1],
ujson.dumps({"test": "test"}))
def main(*argv):
"""
= Configuration import =
check "--help"
Run command: python -m toddler.tools.configimport
:param argv:
:return:
"""
parser = argparse.ArgumentParser(argv, description="ConfigImport v{}".format(__version__))
parser.add_argument("-t", "--type", help="Config type", choices=["crawl"])
if len(argv) > 0:
args = setup(argv, argument_parser=parser, do_not_parse_config=True)
else:
args = setup(argument_parser=parser, do_not_parse_config=True)
print(Style.DIM + Fore.BLUE + "ConfigImport v{}".format(__version__))
with open(args.config) as config_file:
print(Fore.BLUE + "Opened file:" + Fore.RESET + " {}".format(args.config) + Fore.RESET)
if args.type == "crawl":
print(Style.BRIGHT + Fore.BLUE + "Importing crawlConfig")
from toddler.imports.nimbuscrawl import get_configuration
from toddler.models import Host
config_content = config_file.read()
config = get_configuration(config_content)
for host_name, crawl_config in config:
host = Host.objects(host=host_name).first()
if host is None:
host = Host(host=host_name)
host.config["crawlConfig"] = crawl_config
host.save()
print(Fore.GREEN + "+ Added config for host {}".format(host_name))
