diff --git a/tests/flare/datadog-agent-1.tar.bz2 b/tests/flare/datadog-agent-1.tar.bz2 new file mode 100644 index 0000000000..b56f3b974d Binary files /dev/null and b/tests/flare/datadog-agent-1.tar.bz2 differ diff --git a/tests/test_flare.py b/tests/test_flare.py new file mode 100644 index 0000000000..504f7c22a9 --- /dev/null +++ b/tests/test_flare.py @@ -0,0 +1,102 @@ +import unittest +import os.path +import mock +from utils.flare import Flare + +def get_mocked_config(): + return { + 'api_key': 'APIKEY', + 'dd_url': 'https://app.datadoghq.com', + } + +def get_mocked_version(): + return '6.6.6' + +def get_mocked_temp(): + return os.path.join( + os.path.dirname(os.path.realpath(__file__)), + 'flare' + ) + +def mocked_strftime(t): + return '1' + +class FakeResponse(object): + def __init__(self, status_code=200): + self.status_code = status_code + self.text = '{"case_id":1337}' + +class FlareTest(unittest.TestCase): + + @mock.patch('utils.flare.strftime', side_effect=mocked_strftime) + @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp) + @mock.patch('config.get_version', side_effect=get_mocked_version) + @mock.patch('utils.flare.get_config', side_effect=get_mocked_config) + def test_init(self, mock_config, mock_version, mock_tempdir, mock_strftime): + f = Flare(case_id=1337) + conf = mock_config() + self.assertEqual(f._case_id, 1337) + self.assertEqual(f._api_key, conf['api_key']) + self.assertEqual(f._url, 'https://6-6-6-flare.agent.datadoghq.com/support/flare') + self.assertEqual(f._tar_path, os.path.join(get_mocked_temp(), "datadog-agent-1.tar.bz2")) + + @mock.patch('utils.flare.requests.post', return_value=FakeResponse()) + @mock.patch('config.get_version', side_effect=get_mocked_version) + @mock.patch('utils.flare.strftime', side_effect=mocked_strftime) + @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp) + @mock.patch('utils.flare.get_config', side_effect=get_mocked_config) + def test_upload_with_case(self, mock_config, mock_tempdir, mock_stfrtime, mock_version, mock_requests): + f = Flare(case_id=1337) + + assert not mock_requests.called + f.upload(confirmation=False) + assert mock_requests.called + args, kwargs = mock_requests.call_args_list[0] + self.assertEqual( + args, + ('https://6-6-6-flare.agent.datadoghq.com/support/flare/1337?api_key=APIKEY',) + ) + self.assertEqual( + kwargs['files']['flare_file'].name, + os.path.join(get_mocked_temp(), "datadog-agent-1.tar.bz2") + ) + self.assertEqual(kwargs['data']['case_id'], 1337) + self.assertEqual(kwargs['data']['email'], '') + assert kwargs['data']['hostname'] + + @mock.patch('utils.flare.requests.post', return_value=FakeResponse()) + @mock.patch('config.get_version', side_effect=get_mocked_version) + @mock.patch('utils.flare.strftime', side_effect=mocked_strftime) + @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp) + @mock.patch('utils.flare.get_config', side_effect=get_mocked_config) + def test_upload_no_case(self, mock_config, mock_tempdir, mock_stfrtime, mock_version, mock_requests): + f = Flare() + f._ask_for_email = lambda: 'test@example.com' + + assert not mock_requests.called + f.upload(confirmation=False) + assert mock_requests.called + args, kwargs = mock_requests.call_args_list[0] + self.assertEqual( + args, + ('https://6-6-6-flare.agent.datadoghq.com/support/flare?api_key=APIKEY',) + ) + self.assertEqual( + kwargs['files']['flare_file'].name, + os.path.join(get_mocked_temp(), "datadog-agent-1.tar.bz2") + ) + self.assertEqual(kwargs['data']['case_id'], None) + self.assertEqual(kwargs['data']['email'], 'test@example.com') + assert kwargs['data']['hostname'] + + @mock.patch('utils.flare.strftime', side_effect=mocked_strftime) + @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp) + @mock.patch('utils.flare.get_config', side_effect=get_mocked_config) + def test_endpoint(self, mock_config, mock_temp, mock_stfrtime): + f = Flare() + f._ask_for_email = lambda: None + try: + f.upload(confirmation=False) + raise Exception('Should fail before') + except Exception, e: + self.assertEqual(str(e), "Invalid inputs: {'email': None}") diff --git a/utils/flare.py b/utils/flare.py index a8dab2e302..359f2db200 100644 --- a/utils/flare.py +++ b/utils/flare.py @@ -18,7 +18,6 @@ get_config, get_config_path, get_logging_config, - get_os, get_url_endpoint, ) from util import ( @@ -33,9 +32,8 @@ log = logging.getLogger('flare') def configcheck(): - osname = get_os() all_valid = True - for conf_path in glob.glob(os.path.join(get_confd_path(osname), "*.yaml")): + for conf_path in glob.glob(os.path.join(get_confd_path(), "*.yaml")): basename = os.path.basename(conf_path) try: check_yaml(conf_path) @@ -67,6 +65,7 @@ class Flare(object): COMPRESSED_FILE = 'datadog-agent-{0}.tar.bz2' # We limit to 10MB arbitrary MAX_UPLOAD_SIZE = 10485000 + TIMEOUT = 15 def __init__(self, cmdline=False, case_id=None): @@ -123,8 +122,8 @@ def upload(self, confirmation=True): 'hostname': self._hostname, 'email': email } - r = requests.post(url, files=files, data=data) - self._analyse_result(r) + self._resp = requests.post(url, files=files, data=data, timeout=self.TIMEOUT) + self._analyse_result() # Start by creating the tar file which will contain everything def _init_tarfile(self): @@ -303,17 +302,21 @@ def _ask_for_email(self): return raw_input('Please enter your email: ').lower() # Print output (success/error) of the request - def _analyse_result(self, resp): - if resp.status_code in range(200, 203): + def _analyse_result(self): + try: + json_resp = json.loads(self._resp.text) + except ValueError, e: + raise Exception('An unknown error has occured: {0}\n'\ + 'Please contact support by email'.format(self._resp.text)) + if self._resp.status_code in range(200, 203): log.info("Your logs were successfully uploaded. For future reference,"\ - " your internal case id is {0}".format( - json.loads(resp.text)['case_id'])) - elif resp.status_code in range(400, 405): - raise Exception('Your request is incorrect: {0}'.format( - json.loads(resp.text)['error'])) - elif resp.status_code in range(500, 506): + " your internal case id is {0}".format(json_resp['case_id'])) + elif self._resp.status_code in range(400, 405): + raise Exception('Your request is incorrect: {0}'.format(json_resp['error'])) + elif self._resp.status_code in range(500, 506): raise Exception('An error has occurred while uploading: {0}'.format( - json.loads(resp.text)['error'])) + json_resp['error'])) else: - raise Exception('An unknown error has occured: {0}\n'\ - 'Please contact support by email'.format(resp.text)) + raise Exception('An unknown error has occured: {0} - {1}\n'\ + 'Please contact support by email'.format( + self._resp.status_code, self._resp.text))