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
5import atexit
6import base64
7import cgi
8import json
9import logging
10import os
11import requests
12
13LOGGER = logging.getLogger(__name__)
14# Max summaryHtml length (4 KiB) from
15# https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/proto/v1/test_result.proto;drc=ca12b9f52b27f064b0fa47c39baa3b011ffa5790;l=96
16MAX_REPORT_LEN = 4 * 1024
17# VALID_STATUSES is a list of valid status values for test_result['status'].
18# The full list can be obtained at
19# https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/proto/v1/test_result.proto;drc=ca12b9f52b27f064b0fa47c39baa3b011ffa5790;l=151-174
20VALID_STATUSES = {"PASS", "FAIL", "CRASH", "ABORT", "SKIP"}
21
22
23def compose_test_result(test_id, status, expected, test_log=None, tags=None):
24  """Composes the test_result dict item to be posted to result sink.
25
26  Args:
27    test_id: (str) A unique identifier of the test in LUCI context.
28    status: (str) Status of the test. Must be one in |VALID_STATUSES|.
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
34  Returns:
35    A dict of test results with input information, confirming to
36      https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/sink/proto/v1/test_result.proto
37  """
38  assert status in VALID_STATUSES, (
39      '%s is not a valid status (one in %s) for ResultSink.' %
40      (status, VALID_STATUSES))
41
42  for tag in tags or []:
43    assert len(tag) == 2, 'Items in tags should be length 2 tuples of strings'
44    assert isinstance(tag[0], str) and isinstance(
45        tag[1], str), ('Items in'
46                       'tags should be length 2 tuples of strings')
47
48  test_result = {
49      'testId': test_id,
50      'status': status,
51      'expected': expected,
52      'tags': [{
53          'key': key,
54          'value': value
55      } for (key, value) in (tags or [])]
56  }
57
58  if test_log:
59    summary = '<pre>%s</pre>' % cgi.escape(test_log)
60    summary_trunc = ''
61
62    if len(summary) > MAX_REPORT_LEN:
63      summary_trunc = (
64          summary[:MAX_REPORT_LEN - 45] +
65          '...Full output in "Test Log" Artifact.</pre>')
66
67    test_result['summaryHtml'] = summary_trunc or summary
68    if summary_trunc:
69      test_result['artifacts'] = {
70          'Test Log': {
71              'contents': base64.b64encode(test_log)
72          },
73      }
74
75  return test_result
76
77
78class ResultSinkClient(object):
79  """Stores constants and handles posting to ResultSink."""
80
81  def __init__(self):
82    """Initiates and stores constants to class."""
83    self.sink = None
84    luci_context_file = os.environ.get('LUCI_CONTEXT')
85    if not luci_context_file:
86      logging.warning('LUCI_CONTEXT not found in environment. ResultDB'
87                      ' integration disabled.')
88      return
89
90    with open(luci_context_file) as f:
91      self.sink = json.load(f).get('result_sink')
92      if not self.sink:
93        logging.warning('ResultSink constants not found in LUCI context.'
94                        ' ResultDB integration disabled.')
95        return
96
97      self.url = ('http://%s/prpc/luci.resultsink.v1.Sink/ReportTestResults' %
98                  self.sink['address'])
99      self.headers = {
100          'Content-Type': 'application/json',
101          'Accept': 'application/json',
102          'Authorization': 'ResultSink %s' % self.sink['auth_token'],
103      }
104      self._session = requests.Session()
105
106      # Ensure session is closed at exit.
107      atexit.register(self.close)
108
109  def close(self):
110    """Closes the connection to result sink server."""
111    if not self.sink:
112      return
113    LOGGER.info('Closing connection with result sink server.')
114    self._session.close()
115
116  def post(self, test_result):
117    """Posts single test result to server.
118
119    Args:
120        test_result: (dict) Confirming to protocol defined in
121          https://source.chromium.org/chromium/infra/infra/+/master:go/src/go.chromium.org/luci/resultdb/sink/proto/v1/test_result.proto
122    """
123    if not self.sink:
124      return
125
126    res = self._session.post(
127        url=self.url,
128        headers=self.headers,
129        data=json.dumps({'testResults': [test_result]}),
130    )
131    res.raise_for_status()
132