# -*- coding: utf-8 -*-
# SPDX-FileCopyrightText: 2014-2023 Greenbone AG
#
# SPDX-License-Identifier: AGPL-3.0-or-later


# pylint: disable=unused-argument

"""Unit Test for ospd-openvas"""

import logging

from unittest import TestCase
from unittest.mock import patch, MagicMock

from redis.exceptions import ConnectionError as RCE

from ospd.errors import RequiredArgument
from ospd_openvas.db import OpenvasDB, MainDB, ScanDB, KbDB, DBINDEX_NAME, time
from ospd_openvas.errors import OspdOpenvasError

from tests.helper import assert_called


@patch('ospd_openvas.db.redis.Redis')
class TestOpenvasDB(TestCase):
    @patch('ospd_openvas.db.Openvas')
    def test_get_db_connection(
        self, mock_openvas: MagicMock, mock_redis: MagicMock
    ):
        OpenvasDB._db_address = None  # pylint: disable=protected-access
        mock_settings = mock_openvas.get_settings.return_value
        mock_settings.get.return_value = None

        self.assertIsNone(OpenvasDB.get_database_address())

        # set the first time
        mock_openvas.get_settings.return_value = {'db_address': '/foo/bar'}

        self.assertEqual(OpenvasDB.get_database_address(), "unix:///foo/bar")

        self.assertEqual(mock_openvas.get_settings.call_count, 2)

        # should cache address
        self.assertEqual(OpenvasDB.get_database_address(), "unix:///foo/bar")
        self.assertEqual(mock_openvas.get_settings.call_count, 2)

    @patch('ospd_openvas.db.Openvas')
    def test_create_context_fail(self, mock_openvas: MagicMock, mock_redis):
        mock_redis.from_url.side_effect = RCE
        mock_check = mock_openvas.check.return_value
        mock_check.get.return_value = True

        OpenvasDB._db_address = None  # pylint: disable=protected-access
        mock_settings = mock_openvas.get_settings.return_value
        mock_settings.get.return_value = None

        logging.Logger.error = MagicMock()

        with patch.object(time, 'sleep', return_value=None):
            with self.assertRaises(SystemExit):
                OpenvasDB.create_context()

        logging.Logger.error.assert_called_with(  # pylint: disable=no-member
            'Redis Error: Not possible to connect to the kb.'
        )

    @patch('ospd_openvas.db.Openvas')
    def test_create_context_success(self, mock_openvas: MagicMock, mock_redis):
        ctx = mock_redis.from_url.return_value
        mock_check = mock_openvas.check.return_value
        mock_check.get.return_value = True

        OpenvasDB._db_address = None  # pylint: disable=protected-access
        mock_settings = mock_openvas.get_settings.return_value
        mock_settings.get.return_value = None

        ret = OpenvasDB.create_context()
        self.assertIs(ret, ctx)

    def test_select_database_error(self, mock_redis):
        with self.assertRaises(RequiredArgument):
            OpenvasDB.select_database(None, 1)

        with self.assertRaises(RequiredArgument):
            OpenvasDB.select_database(mock_redis, None)

    def test_select_database(self, mock_redis):
        mock_redis.execute_command.return_value = mock_redis

        OpenvasDB.select_database(mock_redis, 1)

        mock_redis.execute_command.assert_called_with('SELECT 1')

    def test_get_list_item_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_list_item(None, 'foo')

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_list_item(ctx, None)

    def test_get_list_item(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.lrange.return_value = ['1234']

        ret = OpenvasDB.get_list_item(ctx, 'name')

        self.assertEqual(ret, ['1234'])
        assert_called(ctx.lrange)

    def test_get_last_list_item(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.rpop.return_value = 'foo'

        ret = OpenvasDB.get_last_list_item(ctx, 'name')

        self.assertEqual(ret, 'foo')
        ctx.rpop.assert_called_with('name')

    def test_get_last_list_item_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_last_list_item(ctx, None)

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_last_list_item(None, 'name')

    def test_remove_list_item(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.lrem.return_value = 1

        OpenvasDB.remove_list_item(ctx, 'name', '1234')

        ctx.lrem.assert_called_once_with('name', count=0, value='1234')

    def test_remove_list_item_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.remove_list_item(None, '1', 'bar')

        with self.assertRaises(RequiredArgument):
            OpenvasDB.remove_list_item(ctx, None, 'bar')

        with self.assertRaises(RequiredArgument):
            OpenvasDB.remove_list_item(ctx, '1', None)

    def test_get_single_item_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_single_item(None, 'foo')

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_single_item(ctx, None)

    def test_get_single_item(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.lindex.return_value = 'a'

        value = OpenvasDB.get_single_item(ctx, 'a')

        self.assertEqual(value, 'a')
        ctx.lindex.assert_called_once_with('a', 0)

    def test_add_single_list(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        pipeline = ctx.pipeline.return_value
        pipeline.delete.return_value = None
        pipeline.execute.return_value = (None, 0)

        OpenvasDB.add_single_list(ctx, 'a', ['12', '11', '12'])

        pipeline.delete.assert_called_once_with('a')
        pipeline.rpush.assert_called_once_with('a', '12', '11', '12')
        assert_called(pipeline.execute)

    def test_add_single_item(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.rpush.return_value = 1

        OpenvasDB.add_single_item(ctx, 'a', ['12', '12'])

        ctx.rpush.assert_called_once_with('a', '12')

    def test_add_single_item_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.add_single_item(None, '1', ['12'])

        with self.assertRaises(RequiredArgument):
            OpenvasDB.add_single_item(ctx, None, ['12'])

        with self.assertRaises(RequiredArgument):
            OpenvasDB.add_single_item(ctx, '1', None)

    def test_set_single_item_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.set_single_item(None, '1', ['12'])

        with self.assertRaises(RequiredArgument):
            OpenvasDB.set_single_item(ctx, None, ['12'])

        with self.assertRaises(RequiredArgument):
            OpenvasDB.set_single_item(ctx, '1', None)

    def test_pop_list_items_no_results(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        pipeline = ctx.pipeline.return_value
        pipeline.lrange.return_value = None
        pipeline.delete.return_value = None
        pipeline.execute.return_value = (None, 0)

        ret = OpenvasDB.pop_list_items(ctx, 'foo')

        self.assertEqual(ret, [])

        pipeline.lrange.assert_called_once_with('foo', 0, -1)
        pipeline.delete.assert_called_once_with('foo')
        assert_called(pipeline.execute)

    def test_pop_list_items_with_results(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        pipeline = ctx.pipeline.return_value
        pipeline.lrange.return_value = None
        pipeline.delete.return_value = None
        pipeline.execute.return_value = [['c', 'b', 'a'], 2]

        ret = OpenvasDB.pop_list_items(ctx, 'results')

        # reversed list
        self.assertEqual(ret, ['a', 'b', 'c'])

        pipeline.lrange.assert_called_once_with('results', 0, -1)
        pipeline.delete.assert_called_once_with('results')
        assert_called(pipeline.execute)

    def test_set_single_item(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        pipeline = ctx.pipeline.return_value
        pipeline.delete.return_value = None
        pipeline.rpush.return_value = None
        pipeline.execute.return_value = None

        OpenvasDB.set_single_item(ctx, 'foo', ['bar'])

        pipeline.delete.assert_called_once_with('foo')
        pipeline.rpush.assert_called_once_with('foo', 'bar')
        assert_called(pipeline.execute)

    def test_get_pattern(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.keys.return_value = ['a', 'b']
        ctx.lrange.return_value = [1, 2, 3]

        ret = OpenvasDB.get_pattern(ctx, 'a')

        self.assertEqual(ret, [['a', [1, 2, 3]], ['b', [1, 2, 3]]])

    def test_get_pattern_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_pattern(None, 'a')

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_pattern(ctx, None)

    def test_get_filenames_and_oids_error(self, mock_redis):
        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_filenames_and_oids(None, None, None)

    def test_get_filenames_and_oids(self, mock_redis):
        def _pars(item):
            return item[4:]

        ctx = mock_redis.from_url.return_value
        ctx.keys.return_value = ['nvt:1', 'nvt:2']
        ctx.lindex.side_effect = ['aa', 'ab']

        ret = OpenvasDB.get_filenames_and_oids(ctx, "nvt:*", _pars)

        self.assertEqual(list(ret), [('aa', '1'), ('ab', '2')])

    def test_get_keys_by_pattern_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_keys_by_pattern(None, 'a')

        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_keys_by_pattern(ctx, None)

    def test_get_keys_by_pattern(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.keys.return_value = ['nvt:2', 'nvt:1']

        ret = OpenvasDB.get_keys_by_pattern(ctx, 'nvt:*')

        # Return sorted list
        self.assertEqual(ret, ['nvt:1', 'nvt:2'])

    def test_get_key_count(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        ctx.keys.return_value = ['aa', 'ab']

        ret = OpenvasDB.get_key_count(ctx, "foo")

        self.assertEqual(ret, 2)
        ctx.keys.assert_called_with('foo')

    def test_get_key_count_with_default_pattern(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        ctx.keys.return_value = ['aa', 'ab']

        ret = OpenvasDB.get_key_count(ctx)

        self.assertEqual(ret, 2)
        ctx.keys.assert_called_with('*')

    def test_get_key_count_error(self, mock_redis):
        with self.assertRaises(RequiredArgument):
            OpenvasDB.get_key_count(None)

    @patch('ospd_openvas.db.Openvas')
    def test_find_database_by_pattern_none(
        self, mock_openvas: MagicMock, mock_redis
    ):
        ctx = mock_redis.from_url.return_value
        ctx.keys.return_value = None

        mock_check = mock_openvas.check.return_value
        mock_check.get.return_value = True

        OpenvasDB._db_address = None  # pylint: disable=protected-access
        mock_settings = mock_openvas.get_settings.return_value
        mock_settings.get.return_value = None

        new_ctx, index = OpenvasDB.find_database_by_pattern('foo*', 123)

        self.assertIsNone(new_ctx)
        self.assertIsNone(index)

    @patch('ospd_openvas.db.Openvas')
    def test_find_database_by_pattern(
        self, mock_openvas: MagicMock, mock_redis
    ):
        ctx = mock_redis.from_url.return_value

        mock_check = mock_openvas.check.return_value
        mock_check.get.return_value = True

        OpenvasDB._db_address = None  # pylint: disable=protected-access
        mock_settings = mock_openvas.get_settings.return_value
        mock_settings.get.return_value = None

        # keys is called twice per iteration
        ctx.keys.side_effect = [None, None, None, None, True, True]

        new_ctx, index = OpenvasDB.find_database_by_pattern('foo*', 123)

        self.assertEqual(new_ctx, ctx)
        self.assertEqual(index, 2)


@patch('ospd_openvas.db.OpenvasDB')
class ScanDBTestCase(TestCase):
    @patch('ospd_openvas.db.redis.Redis')
    def setUp(self, mock_redis):  # pylint: disable=arguments-differ
        self.ctx = mock_redis.from_url.return_value
        self.db = ScanDB(10, self.ctx)

    def test_get_result(self, mock_openvas_db):
        mock_openvas_db.pop_list_items.return_value = [
            'some result',
        ]

        ret = self.db.get_result()

        self.assertEqual(
            ret,
            [
                'some result',
            ],
        )
        mock_openvas_db.pop_list_items.assert_called_with(
            self.ctx, 'internal/results'
        )

    def test_get_status(self, mock_openvas_db):
        mock_openvas_db.get_single_item.return_value = 'some status'

        ret = self.db.get_status('foo')

        self.assertEqual(ret, 'some status')
        mock_openvas_db.get_single_item.assert_called_with(
            self.ctx, 'internal/foo'
        )

    def test_select(self, mock_openvas_db):
        ret = self.db.select(11)

        self.assertIs(ret, self.db)
        self.assertEqual(self.db.index, 11)

        mock_openvas_db.select_database.assert_called_with(self.ctx, 11)

    def test_flush(self, mock_openvas_db):
        self.db.flush()

        self.ctx.flushdb.assert_called_with()


@patch('ospd_openvas.db.OpenvasDB')
class KbDBTestCase(TestCase):
    @patch('ospd_openvas.db.redis.Redis')
    def setUp(self, mock_redis):  # pylint: disable=arguments-differ
        self.ctx = mock_redis.from_url.return_value
        self.db = KbDB(10, self.ctx)

    def test_get_result(self, mock_openvas_db):
        mock_openvas_db.pop_list_items.return_value = [
            'some results',
        ]

        ret = self.db.get_result()

        self.assertEqual(
            ret,
            [
                'some results',
            ],
        )
        mock_openvas_db.pop_list_items.assert_called_with(
            self.ctx, 'internal/results'
        )

    def test_get_status(self, mock_openvas_db):
        mock_openvas_db.get_single_item.return_value = 'some status'

        ret = self.db.get_status('foo')

        self.assertEqual(ret, 'some status')
        mock_openvas_db.get_single_item.assert_called_with(
            self.ctx, 'internal/foo'
        )

    def test_get_scan_status(self, mock_openvas_db):
        status = [
            '192.168.0.1/10/120',
            '192.168.0.2/35/120',
        ]

        mock_openvas_db.pop_list_items.return_value = status

        ret = self.db.get_scan_status()

        self.assertEqual(ret, status)
        mock_openvas_db.pop_list_items.assert_called_with(
            self.ctx, 'internal/status'
        )

    def test_flush(self, mock_openvas_db):
        self.db.flush()

        self.ctx.flushdb.assert_called_with()

    def test_add_scan_id(self, mock_openvas_db):
        self.db.add_scan_id('bar')

        calls = mock_openvas_db.add_single_item.call_args_list

        call = calls[0]
        kwargs = call[0]

        self.assertEqual(kwargs[1], 'internal/bar')
        self.assertEqual(kwargs[2], ['new'])

        call = calls[1]
        kwargs = call[0]

        self.assertEqual(kwargs[1], 'internal/scanid')
        self.assertEqual(kwargs[2], ['bar'])

    def test_add_scan_preferences(self, mock_openvas_db):
        prefs = ['foo', 'bar']

        self.db.add_scan_preferences('foo', prefs)

        mock_openvas_db.add_single_item.assert_called_with(
            self.ctx, 'internal/foo/scanprefs', prefs
        )

    @patch('ospd_openvas.db.OpenvasDB')
    def test_add_credentials_to_scan_preferences(
        self, mock_redis, mock_openvas_db
    ):
        prefs = ['foo', 'bar']

        ctx = mock_redis.from_url.return_value
        mock_openvas_db.create_context.return_value = ctx

        self.db.add_credentials_to_scan_preferences('scan_id', prefs)

        mock_openvas_db.create_context.assert_called_with(
            self.db.index, encoding='utf-8'
        )

        mock_openvas_db.add_single_item.assert_called_with(
            ctx, 'internal/scan_id/scanprefs', prefs
        )

    def test_add_scan_process_id(self, mock_openvas_db):
        self.db.add_scan_process_id(123)

        mock_openvas_db.add_single_item.assert_called_with(
            self.ctx, 'internal/ovas_pid', [123]
        )

    def test_get_scan_process_id(self, mock_openvas_db):
        mock_openvas_db.get_single_item.return_value = '123'

        ret = self.db.get_scan_process_id()

        self.assertEqual(ret, '123')
        mock_openvas_db.get_single_item.assert_called_with(
            self.ctx, 'internal/ovas_pid'
        )

    def test_remove_scan_database(self, mock_openvas_db):
        scan_db = MagicMock(spec=ScanDB)
        scan_db.index = 123

        self.db.remove_scan_database(scan_db)

        mock_openvas_db.remove_list_item.assert_called_with(
            self.ctx, 'internal/dbindex', 123
        )

    def test_target_is_finished_false(self, mock_openvas_db):
        mock_openvas_db.get_single_item.side_effect = ['new']

        ret = self.db.target_is_finished('bar')

        self.assertFalse(ret)

        calls = mock_openvas_db.get_single_item.call_args_list

        call = calls[0]
        args = call[0]

        self.assertEqual(args[1], 'internal/bar')

    def test_target_is_finished_true(self, mock_openvas_db):
        mock_openvas_db.get_single_item.side_effect = ['finished']

        ret = self.db.target_is_finished('bar')

        self.assertTrue(ret)

        calls = mock_openvas_db.get_single_item.call_args_list

        call = calls[0]
        args = call[0]

        self.assertEqual(args[1], 'internal/bar')

    def test_stop_scan(self, mock_openvas_db):
        self.db.stop_scan('foo')

        mock_openvas_db.set_single_item.assert_called_with(
            self.ctx, 'internal/foo', ['stop_all']
        )

    def test_scan_is_stopped_false(self, mock_openvas_db):
        mock_openvas_db.get_single_item.return_value = 'new'

        ret = self.db.scan_is_stopped('foo')

        self.assertFalse(ret)
        mock_openvas_db.get_single_item.assert_called_with(
            self.ctx, 'internal/foo'
        )

    def test_scan_is_stopped_true(self, mock_openvas_db):
        mock_openvas_db.get_single_item.return_value = 'stop_all'

        ret = self.db.scan_is_stopped('foo')

        self.assertTrue(ret)
        mock_openvas_db.get_single_item.assert_called_with(
            self.ctx, 'internal/foo'
        )

    def test_get_scan_databases(self, mock_openvas_db):
        mock_openvas_db.get_list_item.return_value = [
            '4',
            self.db.index,
            '7',
            '11',
        ]

        scan_dbs = self.db.get_scan_databases()

        scan_db = next(scan_dbs)
        self.assertEqual(scan_db.index, '4')

        scan_db = next(scan_dbs)
        self.assertEqual(scan_db.index, '7')

        scan_db = next(scan_dbs)
        self.assertEqual(scan_db.index, '11')

        with self.assertRaises(StopIteration):
            next(scan_dbs)


@patch('ospd_openvas.db.redis.Redis')
class MainDBTestCase(TestCase):
    def test_max_database_index_fail(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.config_get.return_value = {}

        maindb = MainDB(ctx)

        with self.assertRaises(OspdOpenvasError):
            max_db = (  # pylint: disable=unused-variable
                maindb.max_database_index
            )

        ctx.config_get.assert_called_with('databases')

    def test_max_database_index(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.config_get.return_value = {'databases': '123'}

        maindb = MainDB(ctx)

        max_db = maindb.max_database_index

        self.assertEqual(max_db, 123)
        ctx.config_get.assert_called_with('databases')

    def test_try_database_success(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.hsetnx.return_value = 1

        maindb = MainDB(ctx)

        ret = maindb.try_database(1)

        self.assertEqual(ret, True)
        ctx.hsetnx.assert_called_with(DBINDEX_NAME, 1, 1)

    def test_try_database_false(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.hsetnx.return_value = 0

        maindb = MainDB(ctx)

        ret = maindb.try_database(1)

        self.assertEqual(ret, False)
        ctx.hsetnx.assert_called_with(DBINDEX_NAME, 1, 1)

    def test_try_db_index_error(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.hsetnx.side_effect = Exception

        maindb = MainDB(ctx)

        with self.assertRaises(OspdOpenvasError):
            maindb.try_database(1)

    def test_release_database_by_index(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.hdel.return_value = 1

        maindb = MainDB(ctx)

        maindb.release_database_by_index(3)

        ctx.hdel.assert_called_once_with(DBINDEX_NAME, 3)

    def test_release_database(self, mock_redis):
        ctx = mock_redis.from_url.return_value
        ctx.hdel.return_value = 1

        db = MagicMock()
        db.index = 3
        maindb = MainDB(ctx)
        maindb.release_database(db)

        ctx.hdel.assert_called_once_with(DBINDEX_NAME, 3)
        db.flush.assert_called_with()

    def test_release(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        maindb = MainDB(ctx)
        maindb.release()

        ctx.hdel.assert_called_with(DBINDEX_NAME, maindb.index)
        ctx.flushdb.assert_called_with()

    @patch('ospd_openvas.db.Openvas')
    def test_get_new_kb_database(self, mock_openvas: MagicMock, mock_redis):
        ctx = mock_redis.from_url.return_value

        mock_check = mock_openvas.check.return_value
        mock_check.get.return_value = True

        OpenvasDB._db_address = None  # pylint: disable=protected-access
        mock_settings = mock_openvas.get_settings.return_value
        mock_settings.get.return_value = None

        maindb = MainDB(ctx)
        maindb._max_dbindex = 123  # pylint: disable=protected-access

        ctx.hsetnx.side_effect = [0, 0, 1]

        kbdb = maindb.get_new_kb_database()

        self.assertEqual(kbdb.index, 3)
        ctx.flushdb.assert_called_once_with()

    def test_get_new_kb_database_none(self, mock_redis):
        ctx = mock_redis.from_url.return_value

        maindb = MainDB(ctx)
        maindb._max_dbindex = 3  # pylint: disable=protected-access

        ctx.hsetnx.side_effect = [0, 0, 0]

        kbdb = maindb.get_new_kb_database()

        self.assertIsNone(kbdb)
        ctx.flushdb.assert_not_called()

    @patch('ospd_openvas.db.OpenvasDB')
    def test_find_kb_database_by_scan_id_none(
        self, mock_openvas_db, mock_redis
    ):
        ctx = mock_redis.from_url.return_value

        new_ctx = 'bar'  # just some object to compare
        mock_openvas_db.create_context.return_value = new_ctx
        mock_openvas_db.get_key_count.return_value = None

        maindb = MainDB(ctx)
        maindb._max_dbindex = 2  # pylint: disable=protected-access

        kbdb = maindb.find_kb_database_by_scan_id('foo')

        mock_openvas_db.get_key_count.assert_called_once_with(
            new_ctx, 'internal/foo'
        )

        self.assertIsNone(kbdb)

    @patch('ospd_openvas.db.OpenvasDB')
    def test_find_kb_database_by_scan_id(self, mock_openvas_db, mock_redis):
        ctx = mock_redis.from_url.return_value

        new_ctx = 'foo'  # just some object to compare
        mock_openvas_db.create_context.return_value = new_ctx
        mock_openvas_db.get_key_count.side_effect = [0, 1]

        maindb = MainDB(ctx)
        maindb._max_dbindex = 3  # pylint: disable=protected-access

        kbdb = maindb.find_kb_database_by_scan_id('foo')

        mock_openvas_db.get_key_count.assert_called_with(
            new_ctx, 'internal/foo'
        )
        self.assertEqual(kbdb.index, 2)
        self.assertIs(kbdb.ctx, new_ctx)