Skip to content

Commit

Permalink
Refactors test_ingest_events.py removing Mocks
Browse files Browse the repository at this point in the history
  • Loading branch information
Mark Kasaboski committed Nov 29, 2024
1 parent b3bd86d commit d9fc6a3
Show file tree
Hide file tree
Showing 4 changed files with 350 additions and 214 deletions.
10 changes: 8 additions & 2 deletions packages/flare/bin/cron_job_ingest_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
)
154 changes: 154 additions & 0 deletions packages/flare/tests/bin/conftest.py
Original file line number Diff line number Diff line change
@@ -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()
Loading

0 comments on commit d9fc6a3

Please sign in to comment.