diff --git a/packages/flare/bin/cron_job_ingest_events.py b/packages/flare/bin/cron_job_ingest_events.py index bf07b0f..853ca49 100644 --- a/packages/flare/bin/cron_job_ingest_events.py +++ b/packages/flare/bin/cron_job_ingest_events.py @@ -76,7 +76,10 @@ class Application(Protocol): def main( - logger: Logger, storage_passwords: StoragePasswords, kvstore: KVStoreCollections + logger: Logger, + storage_passwords: StoragePasswords, + kvstore: KVStoreCollections, + flare_api_cls: FlareAPI, ) -> None: create_collection(kvstore=kvstore) @@ -107,6 +110,7 @@ def main( ingest_metadata_only=ingest_metadata_only, severities=severities_filter, source_types=source_types_filter, + flare_api_cls=flare_api_cls, ): save_last_fetched(kvstore=kvstore) @@ -311,9 +315,10 @@ def fetch_feed( ingest_metadata_only: bool, severities: list[str], source_types: list[str], + flare_api_cls: FlareAPI, ) -> Iterator[tuple[dict, str]]: try: - flare_api = FlareAPI(api_key=api_key, tenant_id=tenant_id) + flare_api: FlareAPI = flare_api_cls(api_key=api_key, tenant_id=tenant_id) next = get_next(kvstore=kvstore, tenant_id=tenant_id) start_date = get_start_date(kvstore=kvstore) @@ -354,4 +359,5 @@ def get_splunk_service(logger: Logger) -> Service: logger=logger, storage_passwords=app.service.storage_passwords, kvstore=app.service.kvstore, + flare_api_cls=FlareAPI, ) diff --git a/packages/flare/tests/bin/conftest.py b/packages/flare/tests/bin/conftest.py new file mode 100644 index 0000000..80ac3e4 --- /dev/null +++ b/packages/flare/tests/bin/conftest.py @@ -0,0 +1,154 @@ +import json +import os +import pytest +import sys + +from datetime import datetime +from pytest import FixtureRequest +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../bin")) +from constants import KV_COLLECTION_NAME + + +class FakeStoragePassword: + def __init__(self, username: str, clear_password: str) -> None: + self._state = { + "username": username, + "clear_password": clear_password, + } + + @property + def content(self: "FakeStoragePassword") -> "FakeStoragePassword": + return self + + @property + def username(self) -> str: + return self._state["username"] + + @property + def clear_password(self) -> str: + return self._state["clear_password"] + + +class FakeStoragePasswords: + def __init__(self, passwords: List[FakeStoragePassword]) -> None: + self._passwords = passwords + + def list(self) -> List[FakeStoragePassword]: + return self._passwords + + +class FakeKVStoreCollectionData: + def __init__(self) -> None: + self._data: dict[str, str] = {} + + def insert(self, data: str) -> dict[str, str]: + entry = json.loads(data) + self._data[entry["_key"]] = entry["value"] + return entry + + def update(self, id: str, data: str) -> dict[str, str]: + entry = json.loads(data) + self._data[id] = entry["value"] + return entry + + def query(self, **query: dict) -> List[Dict[str, str]]: + return [{"_key": key, "value": value} for key, value in self._data.items()] + + +class FakeKVStoreCollection: + def __init__(self) -> None: + self._data = FakeKVStoreCollectionData() + + @property + def data(self) -> FakeKVStoreCollectionData: + return self._data + + +class FakeKVStoreCollections: + def __init__(self) -> None: + self._collections: dict[str, Any] = {} + + def __getitem__(self, key: str) -> FakeKVStoreCollection: + return self._collections[key] + + def __contains__(self, key: str) -> bool: + return key in self._collections + + def create(self, name: str, fields: dict) -> dict[str, Any]: + self._collections[name] = FakeKVStoreCollection() + return {"headers": {}, "reason": "Created", "status": 200, "body": ""} + + +class FakeLogger: + def __init__(self) -> None: + self.messages: List[str] = [] + + def info(self, message: str) -> None: + self.messages.append(f"INFO: {message}") + + def error(self, message: str) -> None: + self.messages.append(f"ERROR: {message}") + + +class FakeFlareAPI: + def __init__(self, api_key: str, tenant_id: int) -> None: + pass + + def fetch_feed_events( + self, + next: Optional[str], + start_date: Optional[datetime], + ingest_metadata_only: bool, + severities: list[str], + source_types: list[str], + ) -> List[tuple[dict, str]]: + return [ + ( + {"actor": "this guy"}, + "first_next_token", + ), + ( + {"actor": "some other guy"}, + "second_next_token", + ), + ] + + +@pytest.fixture +def storage_passwords(request: FixtureRequest) -> FakeStoragePasswords: + passwords: list[FakeStoragePassword] = [] + data: list[tuple[str, str]] = request.param if hasattr(request, "param") else [] + + if data: + for item in data: + passwords.append( + FakeStoragePassword(username=item[0], clear_password=item[1]) + ) + + return FakeStoragePasswords(passwords=passwords) + + +@pytest.fixture +def kvstore(request: FixtureRequest) -> FakeKVStoreCollections: + kvstore = FakeKVStoreCollections() + data: list[tuple[str, str]] = request.param if hasattr(request, "param") else [] + + if data: + kvstore.create(name=KV_COLLECTION_NAME, fields={}) + for item in data: + kvstore[KV_COLLECTION_NAME].data.insert( + json.dumps({"_key": item[0], "value": item[1]}) + ) + + return kvstore + + +@pytest.fixture +def logger() -> FakeLogger: + return FakeLogger() diff --git a/packages/flare/tests/bin/test_ingest_events.py b/packages/flare/tests/bin/test_ingest_events.py index 9433e2d..520d900 100644 --- a/packages/flare/tests/bin/test_ingest_events.py +++ b/packages/flare/tests/bin/test_ingest_events.py @@ -1,23 +1,21 @@ -import json import os import pytest import sys +from conftest import FakeFlareAPI +from conftest import FakeKVStoreCollections +from conftest import FakeLogger +from conftest import FakeStoragePasswords from datetime import date from datetime import datetime -from datetime import timedelta -from typing import Any -from unittest.mock import MagicMock -from unittest.mock import Mock -from unittest.mock import PropertyMock -from unittest.mock import call -from unittest.mock import patch +from freezegun import freeze_time sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../bin")) from constants import CRON_JOB_THRESHOLD_SINCE_LAST_FETCH from constants import KV_COLLECTION_NAME from constants import CollectionKeys +from constants import PasswordKeys from cron_job_ingest_events import fetch_feed from cron_job_ingest_events import get_api_key from cron_job_ingest_events import get_collection_value @@ -30,198 +28,185 @@ from cron_job_ingest_events import save_last_ingested_tenant_id -def test_get_collection_value_expect_none() -> None: - kvstore = MagicMock() +def test_get_collection_value_expect_none(kvstore: FakeKVStoreCollections) -> None: assert get_collection_value(kvstore=kvstore, key="some_key") is None -def test_get_collection_value_expect_result() -> None: - kvstore = MagicMock() - kvstore.__contains__.side_effect = lambda x: x == KV_COLLECTION_NAME - kvstore[KV_COLLECTION_NAME].data.query.return_value = [ - { - "_key": "some_key", - "value": "some_value", - }, - ] - +@pytest.mark.parametrize("kvstore", [[("some_key", "some_value")]], indirect=True) +def test_get_collection_value_expect_result(kvstore: FakeKVStoreCollections) -> None: assert get_collection_value(kvstore=kvstore, key="some_key") == "some_value" -def test_save_collection_value_expect_insert() -> None: - key = "some_key" - value = "some_value" - kvstore = MagicMock() - save_collection_value(kvstore=kvstore, key=key, value=value) - kvstore[KV_COLLECTION_NAME].data.insert.assert_called_once_with( - json.dumps({"_key": key, "value": value}) - ) +def test_save_collection_value_expect_insert(kvstore: FakeKVStoreCollections) -> None: + save_collection_value(kvstore=kvstore, key="some_key", value="some_value") + assert kvstore[KV_COLLECTION_NAME].data.query() == [ + {"_key": "some_key", "value": "some_value"} + ] -def test_save_collection_value_expect_update() -> None: +@pytest.mark.parametrize("kvstore", [[("some_key", "old_value")]], indirect=True) +def test_save_collection_value_expect_update(kvstore: FakeKVStoreCollections) -> None: key = "some_key" value = "update_value" - kvstore = MagicMock() - kvstore.__contains__.side_effect = lambda x: x == KV_COLLECTION_NAME - kvstore[KV_COLLECTION_NAME].data.query.return_value = [ - { - "_key": key, - "value": "old_value", - }, + + assert kvstore[KV_COLLECTION_NAME].data.query() == [ + {"_key": key, "value": "old_value"} ] save_collection_value(kvstore=kvstore, key=key, value=value) - kvstore[KV_COLLECTION_NAME].data.update.assert_called_once_with( - id=key, - data=json.dumps({"value": value}), - ) + assert kvstore[KV_COLLECTION_NAME].data.query() == [{"_key": key, "value": value}] -def test_get_api_key_tenant_id_expect_exception() -> None: - storage_passwords = MagicMock() - +@pytest.mark.parametrize("storage_passwords", [[]], indirect=True) +def test_get_api_key_expect_exception(storage_passwords: FakeStoragePasswords) -> None: with pytest.raises(Exception, match="API key not found"): get_api_key(storage_passwords=storage_passwords) + +@pytest.mark.parametrize( + "storage_passwords", + [[(PasswordKeys.API_KEY.value, "some_api_key")]], + indirect=True, +) +def test_tenant_id_expect_exception(storage_passwords: FakeStoragePasswords) -> None: with pytest.raises(Exception, match="Tenant ID not found"): get_tenant_id(storage_passwords=storage_passwords) -def test_get_api_credentials_expect_api_key_and_tenant_id() -> None: - storage_passwords = MagicMock() - - api_key_item = Mock() - type(api_key_item.content).username = PropertyMock(return_value="api_key") - type(api_key_item).clear_password = PropertyMock(return_value="some_api_key") - - tenant_id_item = Mock() - type(tenant_id_item.content).username = PropertyMock(return_value="tenant_id") - type(tenant_id_item).clear_password = PropertyMock(return_value=11111) - - storage_passwords.list.return_value = [api_key_item, tenant_id_item] - - api_key = get_api_key(storage_passwords=storage_passwords) - assert api_key == "some_api_key" - tenant_id = get_tenant_id(storage_passwords=storage_passwords) - assert tenant_id == 11111 +@pytest.mark.parametrize( + "storage_passwords", + [ + [ + (PasswordKeys.API_KEY.value, "some_api_key"), + (PasswordKeys.TENANT_ID.value, 11111), + ], + ], + indirect=True, +) +def test_get_api_credentials_expect_api_key_and_tenant_id( + storage_passwords: FakeStoragePasswords, +) -> None: + assert get_api_key(storage_passwords=storage_passwords) == "some_api_key" + assert get_tenant_id(storage_passwords=storage_passwords) == 11111 -@patch( - "cron_job_ingest_events.get_collection_value", return_value="not_an_isoformat_date" +@pytest.mark.parametrize( + "kvstore", + [[(CollectionKeys.START_DATE.value, "non_date_parseable_value")]], + indirect=True, ) -def test_get_start_date_expect_none(get_collection_value_mock: MagicMock) -> None: - kvstore = MagicMock() +def test_get_start_date_expect_none(kvstore: FakeKVStoreCollections) -> None: assert get_start_date(kvstore=kvstore) is None -@patch( - "cron_job_ingest_events.get_collection_value", return_value=date.today().isoformat() +@pytest.mark.parametrize( + "kvstore", + [[(CollectionKeys.START_DATE.value, "2000-01-01")]], + indirect=True, ) -def test_get_start_date_expect_date(get_collection_value_mock: MagicMock) -> None: - kvstore = MagicMock() - assert isinstance(get_start_date(kvstore), date) +def test_get_start_date_expect_date(kvstore: FakeKVStoreCollections) -> None: + assert get_start_date(kvstore) == date(2000, 1, 1) -@patch("cron_job_ingest_events.get_collection_value", return_value="not_a_number") def test_get_last_ingested_tenant_id_expect_none( - get_collection_value_mock: MagicMock, + kvstore: FakeKVStoreCollections, ) -> None: - kvstore = MagicMock() assert get_last_ingested_tenant_id(kvstore=kvstore) is None -@patch("cron_job_ingest_events.get_collection_value", return_value="11111") +@pytest.mark.parametrize( + "kvstore", + [[(CollectionKeys.LAST_INGESTED_TENANT_ID.value, "11111")]], + indirect=True, +) def test_get_last_ingested_tenant_id_expect_integer( - get_collection_value_mock: MagicMock, + kvstore: FakeKVStoreCollections, ) -> None: - kvstore = MagicMock() assert get_last_ingested_tenant_id(kvstore=kvstore) == 11111 -@patch( - "cron_job_ingest_events.get_collection_value", return_value="not_an_isoformat_date" -) -def test_get_last_fetched_expect_none(get_collection_value_mock: MagicMock) -> None: - kvstore = MagicMock() - assert get_last_fetched(kvstore=kvstore) is None - +@freeze_time("2000-01-01") +def test_save_last_ingested_tenant_id_expect_new_tenant_id_and_new_start_date( + kvstore: FakeKVStoreCollections, +) -> None: + # First run will not have data yet. + with pytest.raises(KeyError, match="'event_ingestion_collection'"): + kvstore[KV_COLLECTION_NAME] -@patch( - "cron_job_ingest_events.get_collection_value", - return_value=datetime.now().isoformat(), -) -def test_get_last_fetched_expect_datetime(get_collection_value_mock: MagicMock) -> None: - kvstore = MagicMock() - assert isinstance(get_last_fetched(kvstore=kvstore), datetime) + save_last_ingested_tenant_id(kvstore=kvstore, tenant_id=11111) + assert kvstore[KV_COLLECTION_NAME].data.query() == [ + {"_key": CollectionKeys.START_DATE.value, "value": "2000-01-01"}, + {"_key": CollectionKeys.LAST_INGESTED_TENANT_ID.value, "value": 11111}, + ] -@patch("cron_job_ingest_events.save_collection_value") -@patch("cron_job_ingest_events.get_last_ingested_tenant_id", return_value=None) -def test_save_last_ingested_tenant_id_expect_save_collection_value_called_and_tenant_id_unchanged( - get_last_ingested_tenant_id_mock: MagicMock, - save_collection_value_mock: MagicMock, -) -> None: - kvstore = MagicMock() - save_last_ingested_tenant_id(kvstore=kvstore, tenant_id=11111) - save_collection_value_mock.assert_has_calls( +@pytest.mark.parametrize( + "kvstore", + [ [ - call( - kvstore=kvstore, - key=CollectionKeys.START_DATE.value, - value=date.today().isoformat(), - ), - call( - kvstore=kvstore, - key=CollectionKeys.LAST_INGESTED_TENANT_ID.value, - value=11111, - ), + (CollectionKeys.START_DATE.value, "1999-12-12"), + (CollectionKeys.LAST_INGESTED_TENANT_ID.value, 11111), ] - ) + ], + indirect=True, +) +@freeze_time("2000-01-01") +def test_save_last_ingested_tenant_id_expect_updated_tenant_id_and_updated_start_date( + kvstore: FakeKVStoreCollections, +) -> None: + assert kvstore[KV_COLLECTION_NAME].data.query() == [ + {"_key": CollectionKeys.START_DATE.value, "value": "1999-12-12"}, + {"_key": CollectionKeys.LAST_INGESTED_TENANT_ID.value, "value": 11111}, + ] + save_last_ingested_tenant_id(kvstore=kvstore, tenant_id=22222) + assert kvstore[KV_COLLECTION_NAME].data.query() == [ + {"_key": CollectionKeys.START_DATE.value, "value": "2000-01-01"}, + {"_key": CollectionKeys.LAST_INGESTED_TENANT_ID.value, "value": 22222}, + ] -@patch("cron_job_ingest_events.save_collection_value") -@patch("cron_job_ingest_events.get_start_date", return_value=date.today()) -@patch("cron_job_ingest_events.get_last_ingested_tenant_id", return_value=22222) -def test_save_last_ingested_tenant_id_expect_save_collection_value_not_called_and_tenant_id_changed( - get_last_ingested_tenant_id_mock: MagicMock, - get_start_date_mock: MagicMock, - save_collection_value_mock: MagicMock, -) -> None: - kvstore = MagicMock() - save_last_ingested_tenant_id(kvstore=kvstore, tenant_id=11111) - save_collection_value_mock.assert_has_calls( +@pytest.mark.parametrize( + "kvstore", + [ [ - call( - kvstore=kvstore, - key=CollectionKeys.START_DATE.value, - value=date.today().isoformat(), - ), - call( - kvstore=kvstore, - key=CollectionKeys.LAST_INGESTED_TENANT_ID.value, - value=11111, - ), + (CollectionKeys.START_DATE.value, "1999-12-12"), + (CollectionKeys.LAST_INGESTED_TENANT_ID.value, 11111), ] - ) - - -@patch("cron_job_ingest_events.save_collection_value") -@patch("cron_job_ingest_events.get_start_date", return_value=date.today()) -@patch("cron_job_ingest_events.get_last_ingested_tenant_id", return_value=11111) + ], + indirect=True, +) def test_save_last_ingested_tenant_id_expect_same_tenant_id( - get_last_ingested_tenant_id_mock: MagicMock, - get_start_date_mock: MagicMock, - save_collection_value_mock: MagicMock, + kvstore: FakeKVStoreCollections, ) -> None: - kvstore = MagicMock() save_last_ingested_tenant_id(kvstore=kvstore, tenant_id=11111) - save_collection_value_mock.assert_called_once_with( - kvstore=kvstore, key=CollectionKeys.LAST_INGESTED_TENANT_ID.value, value=11111 - ) + assert kvstore[KV_COLLECTION_NAME].data.query() == [ + {"_key": CollectionKeys.START_DATE.value, "value": "1999-12-12"}, + {"_key": CollectionKeys.LAST_INGESTED_TENANT_ID.value, "value": 11111}, + ] + + +def test_get_last_fetched_expect_none(kvstore: FakeKVStoreCollections) -> None: + assert get_last_fetched(kvstore=kvstore) is None -def test_fetch_feed_expect_exception() -> None: - logger = MagicMock() - kvstore = MagicMock() +@pytest.mark.parametrize( + "kvstore", + [ + [ + ( + CollectionKeys.TIMESTAMP_LAST_FETCH.value, + datetime(2000, 1, 1, 12, 0, 0).isoformat(), + ) + ] + ], + indirect=True, +) +def test_get_last_fetched_expect_datetime(kvstore: FakeKVStoreCollections) -> None: + assert get_last_fetched(kvstore=kvstore) == datetime(2000, 1, 1, 12, 0, 0) + + +def test_fetch_feed_expect_exception( + logger: FakeLogger, kvstore: FakeKVStoreCollections +) -> None: for _ in fetch_feed( logger=logger, kvstore=kvstore, @@ -233,31 +218,21 @@ def test_fetch_feed_expect_exception() -> None: ): pass - logger.error.assert_called_once_with("Exception=Failed to fetch API Token") + assert logger.messages == [ + "INFO: Fetching tenant_id=11111, next=None, start_date=None", + "ERROR: Exception=Failed to fetch API Token", + ] -@patch("cron_job_ingest_events.FlareAPI") -@patch("time.sleep", return_value=None) def test_fetch_feed_expect_feed_response( - sleep: Any, flare_api_mock: MagicMock, capfd: Any + logger: FakeLogger, + kvstore: FakeKVStoreCollections, ) -> None: - logger = MagicMock() - kvstore = MagicMock() - - next = "some_next_value" - first_item = { - "actor": "this guy", - } - second_item = { - "actor": "some other guy", - } + first_item = ({"actor": "this guy"}, "first_next_token") + second_item = ({"actor": "some other guy"}, "second_next_token") expected_items = [first_item, second_item] - flare_api_mock_instance = flare_api_mock.return_value - flare_api_mock_instance.fetch_feed_events.return_value = iter( - [(first_item, next), (second_item, next)] - ) - events: list[dict] = [] + index = 0 for event, next_token in fetch_feed( logger=logger, kvstore=kvstore, @@ -266,64 +241,64 @@ def test_fetch_feed_expect_feed_response( ingest_metadata_only=False, severities=[], source_types=[], + flare_api_cls=FakeFlareAPI, ): - assert next_token == next - events.append(event) + assert event == expected_items[index][0] + assert next_token == expected_items[index][1] + index += 1 - for i in range(len(events)): - assert events[i] == expected_items[i] + assert logger.messages == [ + "INFO: Fetching tenant_id=11111, next=None, start_date=None" + ] -@patch( - "cron_job_ingest_events.get_last_fetched", - return_value=datetime.now() - timedelta(minutes=5), +@pytest.mark.parametrize( + "kvstore", + [ + [ + ( + CollectionKeys.TIMESTAMP_LAST_FETCH.value, + datetime(2000, 1, 1, 12, 0, 0).isoformat(), + ), + ] + ], + indirect=True, ) -def test_main_expect_early_return(get_last_fetched_mock: MagicMock) -> None: - logger = MagicMock() - storage_passwords = MagicMock() - kvstore = MagicMock() - +@freeze_time("2000-01-01 12:09:00") +def test_main_expect_early_return( + logger: FakeLogger, + storage_passwords: FakeStoragePasswords, + kvstore: FakeKVStoreCollections, +) -> None: main(logger=logger, storage_passwords=storage_passwords, kvstore=kvstore) - logger.info.assert_called_once_with( - f"Fetched events less than {int(CRON_JOB_THRESHOLD_SINCE_LAST_FETCH.seconds / 60)} minutes ago, exiting" - ) + logger.messages == [ + f"INFO: Fetched events less than {int(CRON_JOB_THRESHOLD_SINCE_LAST_FETCH.seconds / 60)} minutes ago, exiting" + ] -@patch("cron_job_ingest_events.fetch_feed") -@patch( - "cron_job_ingest_events.get_ingest_metadata_only", - return_value=(False), -) -@patch( - "cron_job_ingest_events.get_tenant_id", - return_value=(111), -) -@patch( - "cron_job_ingest_events.get_api_key", - return_value=("some_api_key"), -) -@patch( - "cron_job_ingest_events.get_last_fetched", - return_value=datetime.now() - timedelta(minutes=10), +@pytest.mark.parametrize( + "storage_passwords", + [ + [ + (PasswordKeys.API_KEY.value, "some_api_key"), + (PasswordKeys.TENANT_ID.value, 11111), + ] + ], + indirect=True, ) +@freeze_time("2000-01-01") def test_main_expect_normal_run( - get_last_fetched_mock: MagicMock, - get_api_key_mock: MagicMock, - get_tenant_id_mock: MagicMock, - get_ingest_metadata_only_mock: MagicMock, - fetch_feed_mock: MagicMock, + logger: FakeLogger, + storage_passwords: FakeStoragePasswords, + kvstore: FakeKVStoreCollections, ) -> None: - logger = MagicMock() - storage_passwords = MagicMock() - kvstore = MagicMock() - - main(logger=logger, storage_passwords=storage_passwords, kvstore=kvstore) - fetch_feed_mock.assert_called_once_with( + main( logger=logger, + storage_passwords=storage_passwords, kvstore=kvstore, - api_key="some_api_key", - tenant_id=111, - ingest_metadata_only=False, - severities=[], - source_types=[], + flare_api_cls=FakeFlareAPI, ) + assert logger.messages == [ + "INFO: Fetching tenant_id=11111, next=None, start_date=FakeDate(2000, 1, 1)", + "INFO: Fetched 2 events", + ] diff --git a/requirements.tools.txt b/requirements.tools.txt index 85605ff..bbb2d03 100644 --- a/requirements.tools.txt +++ b/requirements.tools.txt @@ -3,3 +3,4 @@ mypy==1.12.1 ruff==0.7.0 splunk-appinspect==3.8.0 isort==5.13.2 +freezegun==1.5.1 \ No newline at end of file