result_sink_util.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  1. # Copyright 2020 The Chromium Authors. All rights reserved.
  2. # Use of this source code is governed by a BSD-style license that can be
  3. # found in the LICENSE file.
  4. import atexit
  5. import base64
  6. import cgi
  7. import json
  8. import logging
  9. import os
  10. import requests
  11. import sys
  12. LOGGER = logging.getLogger(__name__)
  13. # VALID_STATUSES is a list of valid status values for test_result['status'].
  14. # The full list can be obtained at
  15. # https://source.chromium.org/chromium/infra/infra/+/main:go/src/go.chromium.org/luci/resultdb/proto/v1/test_result.proto;drc=ca12b9f52b27f064b0fa47c39baa3b011ffa5790;l=151-174
  16. VALID_STATUSES = {"PASS", "FAIL", "CRASH", "ABORT", "SKIP"}
  17. def _compose_test_result(test_id,
  18. status,
  19. expected,
  20. duration=None,
  21. test_log=None,
  22. tags=None,
  23. file_artifacts=None):
  24. """Composes the test_result dict item to be posted to result sink.
  25. Args:
  26. test_id: (str) A unique identifier of the test in LUCI context.
  27. status: (str) Status of the test. Must be one in |VALID_STATUSES|.
  28. duration: (int) Test duration in milliseconds or None if unknown.
  29. expected: (bool) Whether the status is expected.
  30. test_log: (str) Log of the test. Optional.
  31. tags: (list) List of tags. Each item in list should be a length 2 tuple of
  32. string as ("key", "value"). Optional.
  33. file_artifacts: (dict) IDs to abs paths mapping of existing files to
  34. report as artifact.
  35. Returns:
  36. A dict of test results with input information, confirming to
  37. https://source.chromium.org/chromium/infra/infra/+/main:go/src/go.chromium.org/luci/resultdb/sink/proto/v1/test_result.proto
  38. """
  39. tags = tags or []
  40. file_artifacts = file_artifacts or {}
  41. assert status in VALID_STATUSES, (
  42. '%s is not a valid status (one in %s) for ResultSink.' %
  43. (status, VALID_STATUSES))
  44. for tag in tags:
  45. assert len(tag) == 2, 'Items in tags should be length 2 tuples of strings'
  46. assert isinstance(tag[0], str) and isinstance(
  47. tag[1], str), ('Items in'
  48. 'tags should be length 2 tuples of strings')
  49. test_result = {
  50. 'testId': test_id,
  51. 'status': status,
  52. 'expected': expected,
  53. 'tags': [{
  54. 'key': key,
  55. 'value': value
  56. } for (key, value) in tags],
  57. 'testMetadata': {
  58. 'name': test_id,
  59. }
  60. }
  61. test_result['artifacts'] = {
  62. name: {
  63. 'filePath': file_artifacts[name]
  64. } for name in file_artifacts
  65. }
  66. if test_log:
  67. message = ''
  68. if sys.version_info.major < 3:
  69. message = base64.b64encode(test_log)
  70. else:
  71. # Python3 b64encode takes and returns bytes. The result must be
  72. # serializable in order for the eventual json.dumps to succeed
  73. message = base64.b64encode(test_log.encode('utf-8')).decode('utf-8')
  74. test_result['summaryHtml'] = '<text-artifact artifact-id="Test Log" />'
  75. test_result['artifacts'].update({
  76. 'Test Log': {
  77. 'contents': message
  78. },
  79. })
  80. if not test_result['artifacts']:
  81. test_result.pop('artifacts')
  82. if duration:
  83. test_result['duration'] = '%.9fs' % (duration / 1000.0)
  84. return test_result
  85. class ResultSinkClient(object):
  86. """Stores constants and handles posting to ResultSink."""
  87. def __init__(self):
  88. """Initiates and stores constants to class."""
  89. self.sink = None
  90. luci_context_file = os.environ.get('LUCI_CONTEXT')
  91. if not luci_context_file:
  92. logging.warning('LUCI_CONTEXT not found in environment. ResultDB'
  93. ' integration disabled.')
  94. return
  95. with open(luci_context_file) as f:
  96. self.sink = json.load(f).get('result_sink')
  97. if not self.sink:
  98. logging.warning('ResultSink constants not found in LUCI context.'
  99. ' ResultDB integration disabled.')
  100. return
  101. self.url = ('http://%s/prpc/luci.resultsink.v1.Sink/ReportTestResults' %
  102. self.sink['address'])
  103. self.headers = {
  104. 'Content-Type': 'application/json',
  105. 'Accept': 'application/json',
  106. 'Authorization': 'ResultSink %s' % self.sink['auth_token'],
  107. }
  108. self._session = requests.Session()
  109. # Ensure session is closed at exit.
  110. atexit.register(self.close)
  111. logging.getLogger("requests").setLevel(logging.WARNING)
  112. def close(self):
  113. """Closes the connection to result sink server."""
  114. if not self.sink:
  115. return
  116. LOGGER.info('Closing connection with result sink server.')
  117. # Reset to default logging level of test runner scripts.
  118. logging.getLogger("requests").setLevel(logging.DEBUG)
  119. self._session.close()
  120. def post(self, test_id, status, expected, **kwargs):
  121. """Composes and posts a test and status to result sink.
  122. Args:
  123. test_id: (str) A unique identifier of the test in LUCI context.
  124. status: (str) Status of the test. Must be one in |VALID_STATUSES|.
  125. expected: (bool) Whether the status is expected.
  126. **kwargs: Optional keyword args. Namely:
  127. duration: (int) Test duration in milliseconds or None if unknown.
  128. test_log: (str) Log of the test. Optional.
  129. tags: (list) List of tags. Each item in list should be a length 2 tuple
  130. of string as ("key", "value"). Optional.
  131. file_artifacts: (dict) IDs to abs paths mapping of existing files to
  132. report as artifact.
  133. """
  134. if not self.sink:
  135. return
  136. self._post_test_result(
  137. _compose_test_result(test_id, status, expected, **kwargs))
  138. def _post_test_result(self, test_result):
  139. """Posts single test result to server.
  140. This method assumes |self.sink| is not None.
  141. Args:
  142. test_result: (dict) Confirming to protocol defined in
  143. https://source.chromium.org/chromium/infra/infra/+/main:go/src/go.chromium.org/luci/resultdb/sink/proto/v1/test_result.proto
  144. """
  145. res = self._session.post(
  146. url=self.url,
  147. headers=self.headers,
  148. data=json.dumps({'testResults': [test_result]}),
  149. )
  150. res.raise_for_status()