diff options
-rw-r--r-- | api/utils/influx.py | 19 | ||||
-rw-r--r-- | yardstick/tests/unit/apiserver/utils/test_influx.py | 5 |
2 files changed, 14 insertions, 10 deletions
diff --git a/api/utils/influx.py b/api/utils/influx.py index f391ad972..8f3604745 100644 --- a/api/utils/influx.py +++ b/api/utils/influx.py @@ -1,5 +1,6 @@ ############################################################################## # Copyright (c) 2016 Huawei Technologies Co.,Ltd and others. +# Copyright (c) 2019 Intel Corporation # # All rights reserved. This program and the accompanying materials # are made available under the terms of the Apache License, Version 2.0 @@ -22,25 +23,27 @@ from yardstick import dispatcher logger = logging.getLogger(__name__) - -def get_data_db_client(): +def get_data_db_client(db=None): parser = ConfigParser.ConfigParser() try: parser.read(consts.CONF_FILE) - return _get_influxdb_client(parser) + return _get_influxdb_client(parser, db) except ConfigParser.NoOptionError: logger.error('Can not find the key') raise - -def _get_influxdb_client(parser): +def _get_influxdb_client(parser, db=None): if dispatcher.INFLUXDB not in parser.get('DEFAULT', 'dispatcher'): raise exceptions.InfluxDBConfigurationMissing() ip = _get_ip(parser.get('dispatcher_influxdb', 'target')) user = parser.get('dispatcher_influxdb', 'username') password = parser.get('dispatcher_influxdb', 'password') - db_name = parser.get('dispatcher_influxdb', 'db_name') + if db is None: + db_name = parser.get('dispatcher_influxdb', 'db_name') + else: + db_name = db + return influxdb_client.InfluxDBClient(ip, consts.INFLUXDB_PORT, user, password, db_name) @@ -49,9 +52,9 @@ def _get_ip(url): return urlsplit(url).hostname -def query(query_sql): +def query(query_sql, db=None): try: - client = get_data_db_client() + client = get_data_db_client(db) logger.debug('Start to query: %s', query_sql) return list(client.query(query_sql).get_points()) except RuntimeError: diff --git a/yardstick/tests/unit/apiserver/utils/test_influx.py b/yardstick/tests/unit/apiserver/utils/test_influx.py index 6021d35df..3a97ff292 100644 --- a/yardstick/tests/unit/apiserver/utils/test_influx.py +++ b/yardstick/tests/unit/apiserver/utils/test_influx.py @@ -1,5 +1,6 @@ ############################################################################## # Copyright (c) 2016 Huawei Technologies Co.,Ltd and others. +# Copyright (c) 2019 Intel Corporation. # # All rights reserved. This program and the accompanying materials # are made available under the terms of the Apache License, Version 2.0 @@ -29,7 +30,7 @@ class GetDataDbClientTestCase(base.BaseUnitTestCase): self.assertEqual('fake_client', influx.get_data_db_client()) _mock_parser.read.assert_called_once_with(constants.CONF_FILE) - mock_get_client.assert_called_once_with(_mock_parser) + mock_get_client.assert_called_once_with(_mock_parser, None) @mock.patch.object(influx.logger, 'error') @mock.patch.object(influx, '_get_influxdb_client', @@ -46,7 +47,7 @@ class GetDataDbClientTestCase(base.BaseUnitTestCase): influx.get_data_db_client() _mock_parser.read.assert_called_once_with(constants.CONF_FILE) - mock_get_client.assert_called_once_with(_mock_parser) + mock_get_client.assert_called_once_with(_mock_parser, None) class GetIpTestCase(base.BaseUnitTestCase): |