/tests/core/test_flare.py

https://gitlab.com/meetly/dd-agent · Python · 172 lines · 141 code · 28 blank · 3 comment · 2 complexity · 7923584a293d0a9a44ae345a313faf39 MD5 · raw file

  1. # stdlib
  2. import os.path
  3. import unittest
  4. import re
  5. # 3p
  6. import mock
  7. from nose.plugins.attrib import attr
  8. # project
  9. from utils.flare import Flare
  10. def get_mocked_config():
  11. return {
  12. 'api_key': 'APIKEY',
  13. 'dd_url': 'https://app.datadoghq.com',
  14. }
  15. def get_mocked_version():
  16. return '6.6.6'
  17. def get_mocked_temp():
  18. return os.path.join(
  19. os.path.dirname(os.path.realpath(__file__)),
  20. 'fixtures',
  21. 'flare'
  22. )
  23. mock_cfgs = {
  24. 'uri_password' : 'password_uri.yaml',
  25. }
  26. password_tests = {
  27. 'uri_password' : ' - server: mongodb://datadog:V3pZC7ghx1ne82XkyqLnOW36@localhost:27017/admin',
  28. 'uri_password_2' : ' - server: mongodb://datadog:V3!pZC7ghx1ne8#-2XkyqLnOW36!?@localhost:27017/admin',
  29. 'uri_password_expected' : ' - server: mongodb://datadog:********@localhost:27017/admin',
  30. }
  31. def mocked_strftime(t):
  32. return '1'
  33. class FakeResponse(object):
  34. def __init__(self, status_code=200):
  35. self.status_code = status_code
  36. self.text = '{"case_id":1337}'
  37. def json(self):
  38. return {'case_id': 1337}
  39. def raise_for_status(self):
  40. return None
  41. class FlareTest(unittest.TestCase):
  42. @mock.patch('utils.flare.strftime', side_effect=mocked_strftime)
  43. @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp)
  44. @mock.patch('config.get_version', side_effect=get_mocked_version)
  45. @mock.patch('utils.flare.get_config', side_effect=get_mocked_config)
  46. def test_init(self, mock_config, mock_version, mock_tempdir, mock_strftime):
  47. f = Flare(case_id=1337)
  48. conf = mock_config()
  49. self.assertEqual(f._case_id, 1337)
  50. self.assertEqual(f._api_key, conf['api_key'])
  51. self.assertEqual(f._url, 'https://6-6-6-flare.agent.datadoghq.com/support/flare')
  52. self.assertEqual(f.tar_path, os.path.join(get_mocked_temp(), "datadog-agent-1.tar.bz2"))
  53. @mock.patch('utils.flare.requests.post', return_value=FakeResponse())
  54. @mock.patch('config.get_version', side_effect=get_mocked_version)
  55. @mock.patch('utils.flare.strftime', side_effect=mocked_strftime)
  56. @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp)
  57. @mock.patch('utils.flare.get_config', side_effect=get_mocked_config)
  58. def test_upload_with_case(self, mock_config, mock_tempdir, mock_stfrtime, mock_version, mock_requests):
  59. f = Flare(case_id=1337)
  60. f._ask_for_email = lambda: 'test@example.com'
  61. assert not mock_requests.called
  62. f.upload()
  63. assert mock_requests.called
  64. args, kwargs = mock_requests.call_args_list[0]
  65. self.assertEqual(
  66. args,
  67. ('https://6-6-6-flare.agent.datadoghq.com/support/flare/1337?api_key=APIKEY',)
  68. )
  69. self.assertEqual(
  70. kwargs['files']['flare_file'].name,
  71. os.path.join(get_mocked_temp(), "datadog-agent-1.tar.bz2")
  72. )
  73. self.assertEqual(kwargs['data']['case_id'], 1337)
  74. self.assertEqual(kwargs['data']['email'], 'test@example.com')
  75. assert kwargs['data']['hostname']
  76. @mock.patch('utils.flare.requests.post', return_value=FakeResponse())
  77. @mock.patch('config.get_version', side_effect=get_mocked_version)
  78. @mock.patch('utils.flare.strftime', side_effect=mocked_strftime)
  79. @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp)
  80. @mock.patch('utils.flare.get_config', side_effect=get_mocked_config)
  81. def test_upload_no_case(self, mock_config, mock_tempdir, mock_stfrtime, mock_version, mock_requests):
  82. f = Flare()
  83. f._ask_for_email = lambda: 'test@example.com'
  84. assert not mock_requests.called
  85. f.upload()
  86. assert mock_requests.called
  87. args, kwargs = mock_requests.call_args_list[0]
  88. self.assertEqual(
  89. args,
  90. ('https://6-6-6-flare.agent.datadoghq.com/support/flare?api_key=APIKEY',)
  91. )
  92. self.assertEqual(
  93. kwargs['files']['flare_file'].name,
  94. os.path.join(get_mocked_temp(), "datadog-agent-1.tar.bz2")
  95. )
  96. self.assertEqual(kwargs['data']['case_id'], None)
  97. self.assertEqual(kwargs['data']['email'], 'test@example.com')
  98. assert kwargs['data']['hostname']
  99. @attr(requires='core_integration')
  100. @mock.patch('utils.flare.strftime', side_effect=mocked_strftime)
  101. @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp)
  102. @mock.patch('utils.flare.get_config', side_effect=get_mocked_config)
  103. def test_endpoint(self, mock_config, mock_temp, mock_stfrtime):
  104. f = Flare()
  105. f._ask_for_email = lambda: None
  106. try:
  107. f.upload()
  108. raise Exception('Should fail before')
  109. except Exception, e:
  110. self.assertEqual(str(e), "Your request is incorrect: Invalid inputs: 'API key unknown'")
  111. @attr(requires='core_integration')
  112. @mock.patch('utils.flare.strftime', side_effect=mocked_strftime)
  113. @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp)
  114. @mock.patch('utils.flare.get_config', side_effect=get_mocked_config)
  115. def test_uri_password(self, mock_config, mock_tempdir, mock_strftime):
  116. f = Flare()
  117. _, credentials_log = f._strip_credentials(
  118. os.path.join(get_mocked_temp(), mock_cfgs['uri_password']),
  119. f.CHECK_CREDENTIALS
  120. )
  121. self.assertEqual(
  122. credentials_log,
  123. " - this file contains a credential (password in a uri) which has been removed in the collected version"
  124. )
  125. @attr(requires='core_integration')
  126. @mock.patch('utils.flare.strftime', side_effect=mocked_strftime)
  127. @mock.patch('tempfile.gettempdir', side_effect=get_mocked_temp)
  128. @mock.patch('utils.flare.get_config', side_effect=get_mocked_config)
  129. def test_uri_password_regex(self, mock_config, mock_tempdir, mock_strftime):
  130. f = Flare()
  131. password_uri_pattern = filter(
  132. lambda cred_pattern: cred_pattern.label == 'password in a uri',
  133. f.CHECK_CREDENTIALS
  134. ).pop()
  135. line = re.sub(password_uri_pattern.pattern, password_uri_pattern.replacement, password_tests['uri_password'])
  136. self.assertEqual(
  137. line,
  138. password_tests['uri_password_expected']
  139. )
  140. line = re.sub(password_uri_pattern.pattern, password_uri_pattern.replacement, password_tests['uri_password_2'])
  141. self.assertEqual(
  142. line,
  143. password_tests['uri_password_expected']
  144. )