Skip to content

Commit

Permalink
Add User Agent check in ES shipper unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
girodav committed Nov 28, 2023
1 parent 61f4e08 commit baa8f94
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions tests/shippers/test_es.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@
import pytest
from elasticsearch import SerializationError

import share
from share.environment import get_environment
from share.version import version
from shippers import ElasticsearchShipper, JSONSerializer

_now = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
Expand Down Expand Up @@ -48,7 +51,7 @@

class MockTransport(elasticsearch.Transport):
def __init__(self, *args: Any, **kwargs: Any) -> None:
pass
self.kwargs = kwargs


class MockClient(elasticsearch.Elasticsearch):
Expand All @@ -61,13 +64,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
_failures = []


def mock_bulk(client: Any, actions: list[dict[str, Any]], **kwargs: Any) -> tuple[int, list[dict[str, Any]]]:
def mock_bulk(
client: elasticsearch.Elasticsearch, actions: list[dict[str, Any]], **kwargs: Any
) -> tuple[int, list[dict[str, Any]]]:
global _documents
_documents = [actions]
assert client.transport.kwargs["headers"] == {
"User-Agent": share.utils.create_user_agent(esf_version=version, environment=get_environment())
}
return len(actions), []


def mock_bulk_failure(client: Any, actions: list[dict[str, Any]], **kwargs: Any) -> tuple[int, list[dict[str, Any]]]:
def mock_bulk_failure(
client: elasticsearch.Elasticsearch, actions: list[dict[str, Any]], **kwargs: Any
) -> tuple[int, list[dict[str, Any]]]:
global _failures
_failures = list(map(lambda action: {"create": {"_id": action["_id"], "error": "an error"}}, actions))
return len(actions), _failures
Expand Down

0 comments on commit baa8f94

Please sign in to comment.