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', ))