def _setup_es_indexes(_es_client):
"""Sets up ES and makes the client available."""
# Create models in the test index
for search_app in get_search_apps():
# Clean up in case of any aborted test runs
index_name = search_app.es_model.get_target_index_name()
read_alias = search_app.es_model.get_read_alias()
write_alias = search_app.es_model.get_write_alias()
if index_exists(index_name):
delete_index(index_name)
if alias_exists(read_alias):
delete_alias(read_alias)
if alias_exists(write_alias):
delete_alias(write_alias)
# Create indices and aliases
alias_names = (read_alias, write_alias)
create_index(index_name,
search_app.es_model._doc_type.mapping,
alias_names=alias_names)
yield _es_client
for search_app in get_search_apps():
delete_index(search_app.es_model.get_target_index_name())
def _opensearch_session(_opensearch_client):
"""
Session-scoped fixture that creates OpenSearch indexes that persist for the entire test
session.
"""
# Create models in the test index
for search_app in get_search_apps():
# Clean up in case of any aborted test runs
index_name = search_app.search_model.get_target_index_name()
read_alias = search_app.search_model.get_read_alias()
write_alias = search_app.search_model.get_write_alias()
if index_exists(index_name):
delete_index(index_name)
if alias_exists(read_alias):
delete_alias(read_alias)
if alias_exists(write_alias):
delete_alias(write_alias)
# Create indices and aliases
alias_names = (read_alias, write_alias)
create_index(
index_name, search_app.search_model._doc_type.mapping, alias_names=alias_names,
)
yield _opensearch_client
for search_app in get_search_apps():
delete_index(search_app.search_model.get_target_index_name())
def test_migrate_apps(monkeypatch):
"""Test that migrate_apps() migrates the correct apps."""
migrate_app_mock = Mock()
monkeypatch.setattr('datahub.search.migrate.migrate_app', migrate_app_mock)
apps = {app.name for app in list(get_search_apps())[:2]}
migrate_apps(apps)
assert {args[0][0] for args in migrate_app_mock.call_args_list} == apps
def disconnect_delete_search_signal_receivers(es_with_signals):
"""
Fixture that disables signal receivers that delete documents in Elasticsearch.
This is used in tests targeting rollback behaviour. This is because search tests typically
use the synchronous_on_commit fixture, which doesn't model rollback behaviour correctly.
The signal receivers to disable are determined by checking the signal connected to and the
model observed.
"""
disconnected_signal_receivers = []
search_apps = get_search_apps()
for search_app in search_apps:
app_db_model = search_app.queryset.model
for receiver in search_app.get_signal_receivers():
if receiver.signal is post_delete and receiver.sender is app_db_model:
receiver.disconnect()
disconnected_signal_receivers.append(receiver)
yield
# We reconnect the receivers for completeness, though in theory it's not necessary as
# es_with_signals will disconnect them anyway
for receiver in disconnected_signal_receivers:
receiver.connect()
def test_sync_all_models(sync_model_mock):
"""
Test that if --model is not used, all the search apps are synced.
"""
management.call_command(sync_search.Command())
assert sync_model_mock.apply_async.call_count == len(get_search_apps())
def opensearch_collector_context_manager(opensearch, synchronous_on_commit, request):
"""
Slightly lower-level version of opensearch_with_collector.
Function-scoped pytest fixture that:
- ensures OpenSearch is available for the test
- deletes all documents from OpenSearch at the end of the test
- yields a context manager that can be used to collects all model objects saved so
they can be synced to OpenSearch in bulk
Call opensearch_collector_context_manager.flush_and_refresh() to sync collected objects to
OpenSearch and refresh all indices.
In most cases, you should not use this fixture directly, but use opensearch_with_collector or
opensearch_with_signals instead.
"""
marker_apps = {
app
for marker in request.node.iter_markers('opensearch_collector_apps')
for app in marker.args
}
apps = marker_apps or get_search_apps()
yield SavedObjectCollector(opensearch, apps)
def test_reconnects_if_was_connected(self, es_with_signals):
"""Test that signal receivers are reconnected on context manager exit."""
with disable_search_signal_receivers(SimpleModel):
pass
assert all(receiver.is_connected for search_app in get_search_apps()
for receiver in search_app.get_signal_receivers()
if receiver.sender is SimpleModel)
def add_arguments(self, parser):
"""Handle arguments."""
parser.add_argument(
'--model',
action='append',
choices=[search_app.name for search_app in get_search_apps()],
help='Search apps to migrate. If empty, all are migrated.',
)
def test_sync_all_models(monkeypatch):
"""Test that the sync_all_models task starts sub-tasks to sync all models."""
sync_model_mock = Mock()
monkeypatch.setattr('datahub.search.tasks.sync_model', sync_model_mock)
sync_all_models.apply_async()
tasks_created = {call[1]['args'][0] for call in sync_model_mock.apply_async.call_args_list}
assert tasks_created == {app.name for app in get_search_apps()}
def pytest_generate_tests(metafunc):
"""Parametrises tests that use the `search_app` fixture."""
if 'search_app' in metafunc.fixturenames:
apps = get_search_apps()
metafunc.parametrize(
'search_app',
apps,
ids=[app.__class__.__name__ for app in apps],
)
def sync_all_models():
"""
Task that starts sub-tasks to sync all models to Elasticsearch.
acks_late is set to True so that the task restarts if interrupted.
priority is set to the lowest priority (for Redis, 0 is the highest priority).
"""
for search_app in get_search_apps():
sync_model.apply_async(args=(search_app.name, ), )
def test_sync_synchronously(sync_model_mock):
"""
Test that --foreground can be used to run the command in a synchronous (blocking) fashion.
"""
app = get_search_apps()[0]
management.call_command(sync_search.Command(),
model=[app.name],
foreground=True)
assert sync_model_mock.apply.call_count == 1
assert not sync_model_mock.apply_async.called
def __init__(self, *args, **kwargs):
"""Initialises self.entity_by_name dynamically."""
super().__init__(*args, **kwargs)
self.entity_by_name = {
search_app.name: EntitySearch(
search_app.es_model,
search_app.name,
)
for search_app in get_search_apps()
}
def opensearch_with_signals(opensearch, synchronous_on_commit):
"""
Function-scoped pytest fixture that:
- ensures OpenSearch is available for the test
- connects search signal receivers so that OpenSearch documents are automatically
created for model instances saved during the test
- deletes all documents from OpenSearch at the end of the test
Use this fixture when specifically testing search signal receivers.
Call opensearch_with_signals.indices.refresh() after creating objects to refresh all
search indices and ensure synced objects are available for querying.
"""
for search_app in get_search_apps():
search_app.connect_signals()
yield opensearch
for search_app in get_search_apps():
search_app.disconnect_signals()
def setup_es(_setup_es_indexes, synchronous_on_commit):
"""Sets up ES and deletes all the records after each run."""
for search_app in get_search_apps():
search_app.connect_signals()
yield _setup_es_indexes
for search_app in get_search_apps():
search_app.disconnect_signals()
_setup_es_indexes.indices.refresh()
indices = [
search_app.es_model.get_target_index_name()
for search_app in get_search_apps()
]
_setup_es_indexes.delete_by_query(
indices,
body={'query': {
'match_all': {}
}},
)
_setup_es_indexes.indices.refresh()
def test_does_not_reconnect_if_was_disconnected(self):
"""
Test that signal receivers are not reconnected when not originally connected.
Note: Signal receivers are not connected as the es_with_signals fixture is not used.
"""
with disable_search_signal_receivers(SimpleModel):
pass
assert not any(receiver.is_connected
for search_app in get_search_apps()
for receiver in search_app.get_signal_receivers()
if receiver.sender is SimpleModel)
def _get_permission_filters(request):
"""
Gets the permissions filters that should be applied to each search entity (to enforce
permissions).
Only entities that the user has access are returned.
"""
for app in get_search_apps():
if not has_permissions_for_app(request, app):
continue
filter_args = app.get_permission_filters(request)
yield (app.es_model._doc_type.name, filter_args)
def test_sync_model(monkeypatch):
"""Test that the sync_model task starts an OpenSearch sync for that model."""
get_search_app_mock = Mock()
monkeypatch.setattr('datahub.search.tasks.get_search_app', get_search_app_mock)
sync_app_mock = Mock()
monkeypatch.setattr('datahub.search.tasks.sync_app', sync_app_mock)
search_app = next(iter(get_search_apps()))
sync_model.apply_async(args=(search_app.name,))
get_search_app_mock.assert_called_once_with(search_app.name)
sync_app_mock.assert_called_once_with(get_search_app_mock.return_value)
def add_arguments(self, parser):
"""Handle arguments."""
parser.add_argument(
'--model',
action='append',
choices=[search_app.name for search_app in get_search_apps()],
help='Search apps to initialise. If empty, all are initialised.',
)
parser.add_argument(
'--force-update-mapping',
action='store_true',
help=
'Attempts to update the mapping if the index and alias already exist.',
)
def test_reconnects_if_exception_raised(self, es_with_signals):
"""
Test that signal receivers are reconnected if an exception is raised while the
context manager is active.
"""
try:
with disable_search_signal_receivers(SimpleModel):
raise ValueError
except ValueError:
pass
assert all(receiver.is_connected for search_app in get_search_apps()
for receiver in search_app.get_signal_receivers()
if receiver.sender is SimpleModel)
def add_arguments(self, parser):
"""Handle arguments."""
parser.add_argument(
'--batch_size',
type=int,
help=
'Batch size - number of rows processed at a time (defaults to per-model '
'defaults)',
)
parser.add_argument(
'--model',
action='append',
choices=[search_app.name for search_app in get_search_apps()],
help='Search model to import. If empty, it imports all',
)
def add_arguments(self, parser):
"""Handle arguments."""
# TODO: This argument is actually the search app name, not the model name, and
# the argument should therefore be renamed
parser.add_argument(
'--model',
action='append',
choices=[search_app.name for search_app in get_search_apps()],
help='Search model to import. If empty, it imports all',
)
parser.add_argument(
'--foreground',
action='store_true',
help='If specified, the command runs in the foreground without needing Celery '
'running. (By default, it runs asynchronously using Celery.)',
)
def opensearch(_opensearch_session):
"""
Function-scoped pytest fixture that:
- ensures OpenSearch is available for the test
- deletes all documents from OpenSearch at the end of the test.
"""
yield _opensearch_session
_opensearch_session.indices.refresh()
indices = [search_app.search_model.get_target_index_name() for search_app in get_search_apps()]
_opensearch_session.delete_by_query(
indices,
body={'query': {'match_all': {}}},
)
_opensearch_session.indices.refresh()
def __init__(self):
"""Initialises the object."""
self.signal_receivers_to_add = []
self.signal_receivers_to_disable = []
self.deletions = defaultdict(list)
for search_app in get_search_apps():
model = search_app.queryset.model
# set up the receivers to add when collecting the deleted objects
self.signal_receivers_to_add.append(
SignalReceiver(post_delete, model, self._collect), )
# get the existing post/pre_delete receivers that need to be
# disabled in the meantime
for receiver in search_app.get_signal_receivers():
if receiver.signal in (post_delete, pre_delete):
self.signal_receivers_to_disable.append(receiver)
def pytest_generate_tests(metafunc):
"""Parametrises tests that use the `search_app` or `search_view` fixture."""
if 'search_app' in metafunc.fixturenames:
apps = get_search_apps()
metafunc.parametrize(
'search_app',
apps,
ids=[app.__class__.__name__ for app in apps],
)
if 'search_view' in metafunc.fixturenames:
views = [
*v3_view_registry.values(),
*v4_view_registry.values(),
]
metafunc.parametrize(
'search_view',
views,
ids=[view.__class__.__name__ for view in views],
)
def __init__(self, es_client, apps_to_collect):
"""
Initialises the collector.
:param apps_to_collect: the search apps to monitor the `post_save` signal for (and sync
saved objects for when `flush_and_refresh()` is called)
"""
self.collected_apps = set()
self.es_client = es_client
self.signal_receivers_to_connect = [
SignalReceiver(post_save, search_app.queryset.model, self._collect)
for search_app in set(apps_to_collect)
]
# Disconnect all existing search post_save signal receivers (in case they were connected)
self.signal_receivers_to_disable = [
receiver for search_app in get_search_apps()
for receiver in search_app.get_signal_receivers()
if receiver.signal is post_save
]
def disable_search_signal_receivers(model):
"""
Context manager that disables search signals receivers for a particular model.
This disables any signal receivers for that model in all search apps, not just the search
app corresponding to that model.
"""
signal_receivers = [
receiver for search_app in get_search_apps()
for receiver in search_app.get_signal_receivers()
if receiver.sender == model and receiver.is_connected
]
for receiver in signal_receivers:
receiver.disconnect()
try:
yield
finally:
for receiver in signal_receivers:
receiver.connect()
from datahub.search.management.commands import sync_search
@mock.patch(
'datahub.search.apps.index_exists',
mock.Mock(return_value=False),
)
def test_fails_if_index_doesnt_exist():
"""Tests that if the index doesn't exist, sync_search fails."""
with pytest.raises(CommandError):
management.call_command(sync_search.Command())
@pytest.mark.parametrize(
'search_model',
(app.name for app in get_search_apps()),
)
@mock.patch('datahub.search.management.commands.sync_search.sync_model')
@mock.patch(
'datahub.search.apps.index_exists',
mock.Mock(return_value=True),
)
def test_sync_one_model(sync_model_mock, search_model):
"""
Test that --model can be used to specify what we want to sync.
"""
management.call_command(sync_search.Command(), model=[search_model])
assert sync_model_mock.apply_async.call_count == 1
"""Search views URL config."""
from django.urls import path
from datahub.search.apps import get_search_apps
from datahub.search.views import SearchBasicAPIView
urlpatterns = [
path('search', SearchBasicAPIView.as_view(), name='basic'),
]
for search_app in get_search_apps():
if search_app.view:
urlpatterns.append(
path(
f'search/{search_app.name}',
search_app.view.as_view(search_app=search_app),
name=search_app.name,
))
if search_app.export_view:
urlpatterns.append(
path(
f'search/{search_app.name}/export',
search_app.export_view.as_view(search_app=search_app),
name=f'{search_app.name}-export',
))
