osm_lcm/tests/test_prometheus.py

##
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
#
# For those usages not covered by the Apache License, Version 2.0 please
# contact: alfonso.tiernosepulveda@telefonica.com
##

import asynctest
from osm_lcm.prometheus import Prometheus, initial_prometheus_data
from asynctest.mock import Mock
from osm_lcm.data_utils.database.database import Database

__author__ = "Alfonso Tierno <alfonso.tiernosepulveda@telefonica.com>"


class TestPrometheus(asynctest.TestCase):
    async def setUp(self):
        config = {"uri": "http:prometheus:9090", "path": "/etc/prometheus"}
        # Cleanup singleton Database instance
        Database.instance = None

        self.db = Mock(Database({"database": {"driver": "memory"}}).instance.db)
        Database().instance.db = self.db
        self.p = Prometheus(config, worker_id="1", loop=self.loop)

    @asynctest.fail_on(active_handles=True)
    async def test_start(self):
        # test with database empty
        self.db.get_one.return_value = False
        self.p.update = asynctest.CoroutineMock()
        await self.p.start()
        self.db.create.assert_called_once_with("admin", initial_prometheus_data)
        self.p.update.assert_called_once_with()

        # test with database not empty
        self.db.create.reset_mock()
        self.db.get_one.return_value = initial_prometheus_data
        self.p.update.reset_mock()
        await self.p.start()
        self.db.create.assert_not_called()
        self.p.update.assert_called_once_with()

    @asynctest.fail_on(active_handles=True)
    async def test_update(self):
        self.p.PROMETHEUS_LOCKED_TIME = 1
        number_call_set_one = 0

        def _db_set_one(*args, **kwargs):
            # simulated that database is not locked at first call
            nonlocal number_call_set_one

            number_call_set_one += 1
            if number_call_set_one == 1:
                return
            else:
                return {"update": 1}

        def _check_set_one_calls(set_one_calls):
            # check the three calls to database set_one
            self.assertEqual(
                len(set_one_calls),
                3,
                "Not called three times to db.set_one, two blocks, one unblock",
            )
            self.assertIn(
                "admin", set_one_calls[0][0], "db.set_one collection should be admin"
            )
            first_used_time = set_one_calls[0][1]["update_dict"]["_admin.locked_at"]
            second_used_time = set_one_calls[1][1]["update_dict"]["_admin.locked_at"]
            third_used_time = set_one_calls[2][1]["update_dict"]["_admin.locked_at"]
            self.assertTrue(
                first_used_time != 0 and second_used_time != 0,
                "blocking locked_at time must not be 0",
            )
            self.assertGreater(
                second_used_time,
                first_used_time,
                "Every blocking try must contain a new locked_at time",
            )
            self.assertEqual(
                third_used_time, 0, "For unblocking must be set locked_at=0"
            )

        # check add_jobs
        number_call_set_one = 0
        self.db.get_one.return_value = initial_prometheus_data
        self.db.set_one.side_effect = _db_set_one
        self.p.send_data = asynctest.CoroutineMock(return_value=True)
        add_jobs = {"job1": {"job_name": "job1", "nsr_id": "nsr_id"}}
        await self.p.update(add_jobs=add_jobs)
        set_one_calls = self.db.set_one.call_args_list
        _check_set_one_calls(set_one_calls)
        update_dict = set_one_calls[2][1]["update_dict"]
        unset_dict = set_one_calls[2][1]["unset"]
        expected_final_set = {
            "_admin.locked_at": 0,
            "_admin.locked_by": None,
            "_admin.modified_at": set_one_calls[1][1]["update_dict"][
                "_admin.locked_at"
            ],
            "scrape_configs.job1": add_jobs["job1"],
        }
        self.assertEqual(
            update_dict, expected_final_set, "invalid set and unlock values"
        )
        self.assertEqual(unset_dict, None, "invalid unset and unlock values")

        # check remove_jobs
        number_call_set_one = 0
        remove_jobs = ["job1"]
        self.db.set_one.reset_mock()
        await self.p.update(remove_jobs=remove_jobs)
        set_one_calls = self.db.set_one.call_args_list
        _check_set_one_calls(set_one_calls)
        update_dict = set_one_calls[2][1]["update_dict"]
        unset_dict = set_one_calls[2][1]["unset"]
        expected_final_set = {
            "_admin.locked_at": 0,
            "_admin.locked_by": None,
            "_admin.modified_at": set_one_calls[1][1]["update_dict"][
                "_admin.locked_at"
            ],
        }
        self.assertEqual(
            update_dict, expected_final_set, "invalid set and unlock values"
        )
        self.assertEqual(
            unset_dict, {"scrape_configs.job1": None}, "invalid unset and unlock values"
        )

    def test_parse_job(self):
        text_to_parse = """
            # yaml format with jinja2
            key1: "parsing var1='{{ var1 }}'"
            key2: "parsing var2='{{ var2 }}'"
        """
        vars = {"var1": "VAR1", "var2": "VAR2", "var3": "VAR3"}
        expected = {"key1": "parsing var1='VAR1'", "key2": "parsing var2='VAR2'"}
        result = self.p.parse_job(text_to_parse, vars)
        self.assertEqual(result, expected, "Error at jinja2 parse")


if __name__ == "__main__":
    asynctest.main()