Skip to content
This repository has been archived by the owner on Nov 3, 2021. It is now read-only.

Commit

Permalink
Merge pull request #618 from mozilla/improve_unit_test_helpers
Browse files Browse the repository at this point in the history 
Improve unit test suite to connect to resources once
  • Loading branch information
@pwnbus
pwnbus committed Feb 21, 2018
2 parents f300276 + 3e2e95a commit 4cde0c8
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 65 deletions.
93 changes: 93 additions & 0 deletions tests/suite_helper.py
View file
@@ -0,0 +1,93 @@
#!/usr/bin/env python

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# Copyright (c) 2017 Mozilla Corporation

from configlib import getConfig

from kombu import Connection, Queue, Exchange


import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), "../lib"))
from elasticsearch_client import ElasticsearchClient

from utilities.toUTC import toUTC
from utilities.dot_dict import DotDict


# The following functions before the UnitTest class definition
# are a poor man's way to setup resourcing and ensure
# that we only setup a client on time per test suite run
# todo: Fix this by defining an init method when we switch
# from pytest to UnitTest
def parse_config_file():
global CONFIG_FILE_CONTENTS
try:
CONFIG_FILE_CONTENTS
except NameError:
default_config = os.path.join(os.path.dirname(__file__), "config.conf")
options = DotDict()
options.configfile = default_config

options.esservers = list(getConfig('esservers', 'http://localhost:9200', options.configfile).split(','))

options.alertExchange = getConfig('alertexchange', 'alerts', options.configfile)
options.queueName = getConfig('alertqueuename', 'alertBot', options.configfile)
options.alerttopic = getConfig('alerttopic', 'mozdef.*', options.configfile)

options.mquser = getConfig('mquser', 'guest', options.configfile)
options.mqalertserver = getConfig('mqalertserver', 'localhost', options.configfile)
options.mqpassword = getConfig('mqpassword', 'guest', options.configfile)
options.mqport = getConfig('mqport', 5672, options.configfile)
options.mqack = getConfig('mqack', True, options.configfile)
CONFIG_FILE_CONTENTS = options

return CONFIG_FILE_CONTENTS


def parse_mapping_file():
global MAPPING_FILE_CONTENTS
try:
MAPPING_FILE_CONTENTS
except NameError:
default_mapping_file = os.path.join(os.path.dirname(__file__), "../config/defaultMappingTemplate.json")
with open(default_mapping_file) as data_file:
MAPPING_FILE_CONTENTS = data_file.read()
return MAPPING_FILE_CONTENTS


def setup_es_client(options):
global ES_CLIENT
try:
ES_CLIENT
except NameError:
ES_CLIENT = ElasticsearchClient(list('{0}'.format(s) for s in options.esservers))
return ES_CLIENT


def setup_rabbitmq_client(options):
global RABBITMQ_CLIENT
try:
RABBITMQ_CLIENT
except NameError:
mqConnString = 'amqp://{0}:{1}@{2}:{3}//'.format(options.mquser,
options.mqpassword,
options.mqalertserver,
options.mqport)
mqAlertConn = Connection(mqConnString)
alertExchange = Exchange(name=options.alertExchange, type='topic', durable=True, delivery_mode=1)
alertExchange(mqAlertConn).declare()

alertQueue = Queue(options.queueName,
exchange=alertExchange,
routing_key=options.alerttopic,
durable=False,
no_ack=(not options.mqack))
alertQueue(mqAlertConn).declare()

RABBITMQ_CLIENT = mqAlertConn.Consumer(alertQueue, accept=['json'])
return RABBITMQ_CLIENT
79 changes: 14 additions & 65 deletions tests/unit_test_suite.py
View file
Expand Up @@ -5,56 +5,32 @@
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# Copyright (c) 2017 Mozilla Corporation

from datetime import datetime, timedelta
from dateutil.parser import parse

import random
import pytest

import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), "../lib"))
from elasticsearch_client import ElasticsearchClient

from utilities.toUTC import toUTC
from utilities.dot_dict import DotDict

from datetime import datetime, timedelta
from dateutil.parser import parse

from configlib import getConfig

from kombu import Connection, Queue, Exchange

import random
import pytest
from suite_helper import parse_config_file, parse_mapping_file, setup_es_client, setup_rabbitmq_client


class UnitTestSuite(object):

def setup(self):
self.options = parse_config_file()
self.mapping_options = parse_mapping_file()
self.es_client = setup_es_client(self.options)
self.rabbitmq_alerts_consumer = setup_rabbitmq_client(self.options)

current_date = datetime.now()
self.event_index_name = current_date.strftime("events-%Y%m%d")
self.previous_event_index_name = (current_date - timedelta(days=1)).strftime("events-%Y%m%d")
self.alert_index_name = current_date.strftime("alerts-%Y%m")
self.parse_config()

# Elasticsearch
self.es_client = ElasticsearchClient(list('{0}'.format(s) for s in self.options.esservers))

# RabbitMQ
mqConnString = 'amqp://{0}:{1}@{2}:{3}//'.format(self.options.mquser,
self.options.mqpassword,
self.options.mqalertserver,
self.options.mqport)

mqAlertConn = Connection(mqConnString)
alertExchange = Exchange(name=self.options.alertExchange, type='topic', durable=True, delivery_mode=1)
alertExchange(mqAlertConn).declare()

alertQueue = Queue(self.options.queueName,
exchange=alertExchange,
routing_key=self.options.alerttopic,
durable=False,
no_ack=(not self.options.mqack))
alertQueue(mqAlertConn).declare()

self.rabbitmq_alerts_consumer = mqAlertConn.Consumer(alertQueue, accept=['json'])

if pytest.config.option.delete_indexes:
self.reset_elasticsearch()
Expand All @@ -63,25 +39,6 @@ def setup(self):
if pytest.config.option.delete_queues:
self.reset_rabbitmq()

def parse_config(self):
default_config = os.path.join(os.path.dirname(__file__), "config.conf")
options = DotDict()
options.configfile = default_config

options.esservers = list(getConfig('esservers', 'http://localhost:9200', options.configfile).split(','))

options.alertExchange = getConfig('alertexchange', 'alerts', options.configfile)
options.queueName = getConfig('alertqueuename', 'alertBot', options.configfile)
options.alerttopic = getConfig('alerttopic', 'mozdef.*', options.configfile)

options.mquser = getConfig('mquser', 'guest', options.configfile)
options.mqalertserver = getConfig('mqalertserver', 'localhost', options.configfile)
options.mqpassword = getConfig('mqpassword', 'guest', options.configfile)
options.mqport = getConfig('mqport', 5672, options.configfile)
options.mqack = getConfig('mqack', True, options.configfile)

self.options = options

def reset_rabbitmq(self):
self.rabbitmq_alerts_consumer.channel.queue_purge()

Expand All @@ -91,26 +48,18 @@ def teardown(self):
if pytest.config.option.delete_queues:
self.reset_rabbitmq()

self.rabbitmq_alerts_consumer.connection.close()
self.rabbitmq_alerts_consumer.close()

def populate_test_event(self, event, event_type='event'):
self.es_client.save_event(body=event, doc_type=event_type)

def populate_test_object(self, event, event_type='event'):
self.es_client.save_object(index='events', body=event, doc_type=event_type)

def setup_elasticsearch(self):
default_mapping_file = os.path.join(os.path.dirname(__file__), "../config/defaultMappingTemplate.json")
mapping_str = ''
with open(default_mapping_file) as data_file:
mapping_str = data_file.read()

self.es_client.create_index(self.event_index_name, index_config=mapping_str)
self.es_client.create_index(self.event_index_name, index_config=self.mapping_options)
self.es_client.create_alias('events', self.event_index_name)
self.es_client.create_index(self.previous_event_index_name, index_config=mapping_str)
self.es_client.create_index(self.previous_event_index_name, index_config=self.mapping_options)
self.es_client.create_alias('events-previous', self.previous_event_index_name)
self.es_client.create_index(self.alert_index_name, index_config=mapping_str)
self.es_client.create_index(self.alert_index_name, index_config=self.mapping_options)
self.es_client.create_alias('alerts', self.alert_index_name)

def reset_elasticsearch(self):
Expand Down

0 comments on commit 4cde0c8

Please sign in to comment.