suggest/
store.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
4 */
5
6use std::{
7    collections::{hash_map::Entry, HashMap, HashSet},
8    path::{Path, PathBuf},
9    sync::Arc,
10};
11
12use error_support::{breadcrumb, handle_error, trace};
13use once_cell::sync::OnceCell;
14use parking_lot::Mutex;
15use remote_settings::{self, RemoteSettingsError, RemoteSettingsServer, RemoteSettingsService};
16
17use serde::de::DeserializeOwned;
18
19use crate::{
20    config::{SuggestGlobalConfig, SuggestProviderConfig},
21    db::{ConnectionType, IngestedRecord, Sqlite3Extension, SuggestDao, SuggestDb},
22    error::Error,
23    geoname::{Geoname, GeonameAlternates, GeonameMatch},
24    metrics::{MetricsContext, SuggestIngestionMetrics, SuggestQueryMetrics},
25    provider::{SuggestionProvider, SuggestionProviderConstraints, DEFAULT_INGEST_PROVIDERS},
26    rs::{
27        Client, Collection, DownloadedDynamicRecord, Record, SuggestAttachment, SuggestRecord,
28        SuggestRecordId, SuggestRecordType, SuggestRemoteSettingsClient,
29    },
30    QueryWithMetricsResult, Result, SuggestApiResult, Suggestion, SuggestionQuery,
31};
32
33/// Builder for [SuggestStore]
34///
35/// Using a builder is preferred to calling the constructor directly since it's harder to confuse
36/// the data_path and cache_path strings.
37#[derive(uniffi::Object)]
38pub struct SuggestStoreBuilder(Mutex<SuggestStoreBuilderInner>);
39
40#[derive(Default)]
41struct SuggestStoreBuilderInner {
42    data_path: Option<String>,
43    remote_settings_server: Option<RemoteSettingsServer>,
44    remote_settings_service: Option<Arc<RemoteSettingsService>>,
45    remote_settings_bucket_name: Option<String>,
46    extensions_to_load: Vec<Sqlite3Extension>,
47}
48
49impl Default for SuggestStoreBuilder {
50    fn default() -> Self {
51        Self::new()
52    }
53}
54
55#[uniffi::export]
56impl SuggestStoreBuilder {
57    #[uniffi::constructor]
58    pub fn new() -> SuggestStoreBuilder {
59        Self(Mutex::new(SuggestStoreBuilderInner::default()))
60    }
61
62    pub fn data_path(self: Arc<Self>, path: String) -> Arc<Self> {
63        self.0.lock().data_path = Some(path);
64        self
65    }
66
67    /// Deprecated: this is no longer used by the suggest component.
68    pub fn cache_path(self: Arc<Self>, _path: String) -> Arc<Self> {
69        // We used to use this, but we're not using it anymore, just ignore the call
70        self
71    }
72
73    pub fn remote_settings_server(self: Arc<Self>, server: RemoteSettingsServer) -> Arc<Self> {
74        self.0.lock().remote_settings_server = Some(server);
75        self
76    }
77
78    pub fn remote_settings_bucket_name(self: Arc<Self>, bucket_name: String) -> Arc<Self> {
79        self.0.lock().remote_settings_bucket_name = Some(bucket_name);
80        self
81    }
82
83    pub fn remote_settings_service(
84        self: Arc<Self>,
85        rs_service: Arc<RemoteSettingsService>,
86    ) -> Arc<Self> {
87        self.0.lock().remote_settings_service = Some(rs_service);
88        self
89    }
90
91    /// Add an sqlite3 extension to load
92    ///
93    /// library_name should be the name of the library without any extension, for example `libmozsqlite3`.
94    /// entrypoint should be the entry point, for example `sqlite3_fts5_init`.  If `null` (the default)
95    /// entry point will be used (see https://sqlite.org/loadext.html for details).
96    pub fn load_extension(
97        self: Arc<Self>,
98        library: String,
99        entry_point: Option<String>,
100    ) -> Arc<Self> {
101        self.0.lock().extensions_to_load.push(Sqlite3Extension {
102            library,
103            entry_point,
104        });
105        self
106    }
107
108    #[handle_error(Error)]
109    pub fn build(&self) -> SuggestApiResult<Arc<SuggestStore>> {
110        let inner = self.0.lock();
111        let extensions_to_load = inner.extensions_to_load.clone();
112        let data_path = inner
113            .data_path
114            .clone()
115            .ok_or_else(|| Error::SuggestStoreBuilder("data_path not specified".to_owned()))?;
116        let rs_service = inner.remote_settings_service.clone().ok_or_else(|| {
117            Error::RemoteSettings(RemoteSettingsError::Other {
118                reason: "remote_settings_service_not_specified".to_string(),
119            })
120        })?;
121        Ok(Arc::new(SuggestStore {
122            inner: SuggestStoreInner::new(
123                data_path,
124                extensions_to_load,
125                SuggestRemoteSettingsClient::new(&rs_service),
126            ),
127        }))
128    }
129}
130
131/// What should be interrupted when [SuggestStore::interrupt] is called?
132#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, uniffi::Enum)]
133pub enum InterruptKind {
134    /// Interrupt read operations like [SuggestStore::query]
135    Read,
136    /// Interrupt write operations.  This mostly means [SuggestStore::ingest], but
137    /// other operations may also be interrupted.
138    Write,
139    /// Interrupt both read and write operations,
140    ReadWrite,
141}
142
143/// The store is the entry point to the Suggest component. It incrementally
144/// downloads suggestions from the Remote Settings service, stores them in a
145/// local database, and returns them in response to user queries.
146///
147/// Your application should create a single store, and manage it as a singleton.
148/// The store is thread-safe, and supports concurrent queries and ingests. We
149/// expect that your application will call [`SuggestStore::query()`] to show
150/// suggestions as the user types into the address bar, and periodically call
151/// [`SuggestStore::ingest()`] in the background to update the database with
152/// new suggestions from Remote Settings.
153///
154/// For responsiveness, we recommend always calling `query()` on a worker
155/// thread. When the user types new input into the address bar, call
156/// [`SuggestStore::interrupt()`] on the main thread to cancel the query
157/// for the old input, and unblock the worker thread for the new query.
158///
159/// The store keeps track of the state needed to support incremental ingestion,
160/// but doesn't schedule the ingestion work itself, or decide how many
161/// suggestions to ingest at once. This is for two reasons:
162///
163/// 1. The primitives for scheduling background work vary between platforms, and
164///    aren't available to the lower-level Rust layer. You might use an idle
165///    timer on Desktop, `WorkManager` on Android, or `BGTaskScheduler` on iOS.
166/// 2. Ingestion constraints can change, depending on the platform and the needs
167///    of your application. A mobile device on a metered connection might want
168///    to request a small subset of the Suggest data and download the rest
169///    later, while a desktop on a fast link might download the entire dataset
170///    on the first launch.
171#[derive(uniffi::Object)]
172pub struct SuggestStore {
173    inner: SuggestStoreInner<SuggestRemoteSettingsClient>,
174}
175
176#[uniffi::export]
177impl SuggestStore {
178    /// Creates a Suggest store.
179    #[uniffi::constructor()]
180    pub fn new(path: &str, remote_settings_service: Arc<RemoteSettingsService>) -> Self {
181        let client = SuggestRemoteSettingsClient::new(&remote_settings_service);
182        Self {
183            inner: SuggestStoreInner::new(path.to_owned(), vec![], client),
184        }
185    }
186
187    /// Queries the database for suggestions.
188    #[handle_error(Error)]
189    pub fn query(&self, query: SuggestionQuery) -> SuggestApiResult<Vec<Suggestion>> {
190        Ok(self.inner.query(query)?.suggestions)
191    }
192
193    /// Queries the database for suggestions.
194    #[handle_error(Error)]
195    pub fn query_with_metrics(
196        &self,
197        query: SuggestionQuery,
198    ) -> SuggestApiResult<QueryWithMetricsResult> {
199        self.inner.query(query)
200    }
201
202    /// Dismiss a suggestion.
203    ///
204    /// Dismissed suggestions cannot be fetched again.
205    #[handle_error(Error)]
206    pub fn dismiss_by_suggestion(&self, suggestion: &Suggestion) -> SuggestApiResult<()> {
207        self.inner.dismiss_by_suggestion(suggestion)
208    }
209
210    /// Dismiss a suggestion by its dismissal key.
211    ///
212    /// Dismissed suggestions cannot be fetched again.
213    ///
214    /// Prefer [SuggestStore::dismiss_by_suggestion] if you have a
215    /// `crate::Suggestion`. This method is intended for cases where a
216    /// suggestion originates outside this component.
217    #[handle_error(Error)]
218    pub fn dismiss_by_key(&self, key: &str) -> SuggestApiResult<()> {
219        self.inner.dismiss_by_key(key)
220    }
221
222    /// Deprecated, use [SuggestStore::dismiss_by_suggestion] or
223    /// [SuggestStore::dismiss_by_key] instead.
224    ///
225    /// Dismiss a suggestion
226    ///
227    /// Dismissed suggestions will not be returned again
228    #[handle_error(Error)]
229    pub fn dismiss_suggestion(&self, suggestion_url: String) -> SuggestApiResult<()> {
230        self.inner.dismiss_suggestion(suggestion_url)
231    }
232
233    /// Clear dismissed suggestions
234    #[handle_error(Error)]
235    pub fn clear_dismissed_suggestions(&self) -> SuggestApiResult<()> {
236        self.inner.clear_dismissed_suggestions()
237    }
238
239    /// Return whether a suggestion has been dismissed.
240    ///
241    /// [SuggestStore::query] will never return dismissed suggestions, so
242    /// normally you never need to know whether a `Suggestion` has been
243    /// dismissed, but this method can be used to do so.
244    #[handle_error(Error)]
245    pub fn is_dismissed_by_suggestion(&self, suggestion: &Suggestion) -> SuggestApiResult<bool> {
246        self.inner.is_dismissed_by_suggestion(suggestion)
247    }
248
249    /// Return whether a suggestion has been dismissed given its dismissal key.
250    ///
251    /// [SuggestStore::query] will never return dismissed suggestions, so
252    /// normally you never need to know whether a suggestion has been dismissed.
253    /// This method is intended for cases where a dismissal key originates
254    /// outside this component.
255    #[handle_error(Error)]
256    pub fn is_dismissed_by_key(&self, key: &str) -> SuggestApiResult<bool> {
257        self.inner.is_dismissed_by_key(key)
258    }
259
260    /// Return whether any suggestions have been dismissed.
261    #[handle_error(Error)]
262    pub fn any_dismissed_suggestions(&self) -> SuggestApiResult<bool> {
263        self.inner.any_dismissed_suggestions()
264    }
265
266    /// Interrupts any ongoing queries.
267    ///
268    /// This should be called when the user types new input into the address
269    /// bar, to ensure that they see fresh suggestions as they type. This
270    /// method does not interrupt any ongoing ingests.
271    #[uniffi::method(default(kind = None))]
272    pub fn interrupt(&self, kind: Option<InterruptKind>) {
273        self.inner.interrupt(kind)
274    }
275
276    /// Ingests new suggestions from Remote Settings.
277    #[handle_error(Error)]
278    pub fn ingest(
279        &self,
280        constraints: SuggestIngestionConstraints,
281    ) -> SuggestApiResult<SuggestIngestionMetrics> {
282        self.inner.ingest(constraints)
283    }
284
285    /// Removes all content from the database.
286    #[handle_error(Error)]
287    pub fn clear(&self) -> SuggestApiResult<()> {
288        self.inner.clear()
289    }
290
291    /// Returns global Suggest configuration data.
292    #[handle_error(Error)]
293    pub fn fetch_global_config(&self) -> SuggestApiResult<SuggestGlobalConfig> {
294        self.inner.fetch_global_config()
295    }
296
297    /// Returns per-provider Suggest configuration data.
298    #[handle_error(Error)]
299    pub fn fetch_provider_config(
300        &self,
301        provider: SuggestionProvider,
302    ) -> SuggestApiResult<Option<SuggestProviderConfig>> {
303        self.inner.fetch_provider_config(provider)
304    }
305
306    /// Fetches geonames stored in the database. A geoname represents a
307    /// geographic place.
308    ///
309    /// See `fetch_geonames` in `geoname.rs` for documentation.
310    #[handle_error(Error)]
311    pub fn fetch_geonames(
312        &self,
313        query: &str,
314        match_name_prefix: bool,
315        filter: Option<Vec<Geoname>>,
316    ) -> SuggestApiResult<Vec<GeonameMatch>> {
317        self.inner.fetch_geonames(query, match_name_prefix, filter)
318    }
319
320    /// Fetches a geoname's names stored in the database.
321    ///
322    /// See `fetch_geoname_alternates` in `geoname.rs` for documentation.
323    #[handle_error(Error)]
324    pub fn fetch_geoname_alternates(
325        &self,
326        geoname: &Geoname,
327    ) -> SuggestApiResult<GeonameAlternates> {
328        self.inner.fetch_geoname_alternates(geoname)
329    }
330}
331
332impl SuggestStore {
333    pub fn force_reingest(&self) {
334        self.inner.force_reingest()
335    }
336}
337
338#[cfg(feature = "benchmark_api")]
339impl SuggestStore {
340    /// Creates a WAL checkpoint. This will cause changes in the write-ahead log
341    /// to be written to the DB. See:
342    /// https://sqlite.org/pragma.html#pragma_wal_checkpoint
343    pub fn checkpoint(&self) {
344        self.inner.checkpoint();
345    }
346}
347
348/// Constraints limit which suggestions to ingest from Remote Settings.
349#[derive(Clone, Default, Debug, uniffi::Record)]
350pub struct SuggestIngestionConstraints {
351    #[uniffi(default = None)]
352    pub providers: Option<Vec<SuggestionProvider>>,
353    #[uniffi(default = None)]
354    pub provider_constraints: Option<SuggestionProviderConstraints>,
355    /// Only run ingestion if the table `suggestions` is empty
356    ///
357    // This is indented to handle periodic updates.  Consumers can schedule an ingest with
358    // `empty_only=true` on startup and a regular ingest with `empty_only=false` to run on a long periodic schedule (maybe
359    // once a day). This allows ingestion to normally be run at a slow, periodic rate.  However, if
360    // there is a schema upgrade that causes the database to be thrown away, then the
361    // `empty_only=true` ingestion that runs on startup will repopulate it.
362    #[uniffi(default = false)]
363    pub empty_only: bool,
364}
365
366impl SuggestIngestionConstraints {
367    pub fn all_providers() -> Self {
368        Self {
369            providers: Some(vec![
370                SuggestionProvider::Amp,
371                SuggestionProvider::Wikipedia,
372                SuggestionProvider::Amo,
373                SuggestionProvider::Yelp,
374                SuggestionProvider::Mdn,
375                SuggestionProvider::Weather,
376                SuggestionProvider::Fakespot,
377                SuggestionProvider::Dynamic,
378            ]),
379            ..Self::default()
380        }
381    }
382
383    fn matches_dynamic_record(&self, record: &DownloadedDynamicRecord) -> bool {
384        match self
385            .provider_constraints
386            .as_ref()
387            .and_then(|c| c.dynamic_suggestion_types.as_ref())
388        {
389            None => false,
390            Some(suggestion_types) => suggestion_types.contains(&record.suggestion_type),
391        }
392    }
393
394    fn amp_matching_uses_fts(&self) -> bool {
395        self.provider_constraints
396            .as_ref()
397            .and_then(|c| c.amp_alternative_matching.as_ref())
398            .map(|constraints| constraints.uses_fts())
399            .unwrap_or(false)
400    }
401}
402
403/// The implementation of the store. This is generic over the Remote Settings
404/// client, and is split out from the concrete [`SuggestStore`] for testing
405/// with a mock client.
406pub(crate) struct SuggestStoreInner<S> {
407    /// Path to the persistent SQL database.
408    ///
409    /// This stores things that should persist when the user clears their cache.
410    /// It's not currently used because not all consumers pass this in yet.
411    #[allow(unused)]
412    data_path: PathBuf,
413    dbs: OnceCell<SuggestStoreDbs>,
414    extensions_to_load: Vec<Sqlite3Extension>,
415    settings_client: S,
416}
417
418impl<S> SuggestStoreInner<S> {
419    pub fn new(
420        data_path: impl Into<PathBuf>,
421        extensions_to_load: Vec<Sqlite3Extension>,
422        settings_client: S,
423    ) -> Self {
424        Self {
425            data_path: data_path.into(),
426            extensions_to_load,
427            dbs: OnceCell::new(),
428            settings_client,
429        }
430    }
431
432    /// Returns this store's database connections, initializing them if
433    /// they're not already open.
434    fn dbs(&self) -> Result<&SuggestStoreDbs> {
435        self.dbs
436            .get_or_try_init(|| SuggestStoreDbs::open(&self.data_path, &self.extensions_to_load))
437    }
438
439    fn query(&self, query: SuggestionQuery) -> Result<QueryWithMetricsResult> {
440        let mut metrics = SuggestQueryMetrics::default();
441        let mut suggestions = vec![];
442
443        let unique_providers = query.providers.iter().collect::<HashSet<_>>();
444        let reader = &self.dbs()?.reader;
445        for provider in unique_providers {
446            let new_suggestions = metrics.measure_query(provider.to_string(), || {
447                reader.read(|dao| match provider {
448                    SuggestionProvider::Amp => dao.fetch_amp_suggestions(&query),
449                    SuggestionProvider::Wikipedia => dao.fetch_wikipedia_suggestions(&query),
450                    SuggestionProvider::Amo => dao.fetch_amo_suggestions(&query),
451                    SuggestionProvider::Yelp => dao.fetch_yelp_suggestions(&query),
452                    SuggestionProvider::Mdn => dao.fetch_mdn_suggestions(&query),
453                    SuggestionProvider::Weather => dao.fetch_weather_suggestions(&query),
454                    SuggestionProvider::Fakespot => dao.fetch_fakespot_suggestions(&query),
455                    SuggestionProvider::Dynamic => dao.fetch_dynamic_suggestions(&query),
456                })
457            })?;
458            suggestions.extend(new_suggestions);
459        }
460
461        // Note: it's important that this is a stable sort to keep the intra-provider order stable.
462        // For example, we can return multiple fakespot-suggestions all with `score=0.245`.  In
463        // that case, they must be in the same order that `fetch_fakespot_suggestions` returned
464        // them in.
465        suggestions.sort();
466        if let Some(limit) = query.limit.and_then(|limit| usize::try_from(limit).ok()) {
467            suggestions.truncate(limit);
468        }
469        Ok(QueryWithMetricsResult {
470            suggestions,
471            query_times: metrics.times,
472        })
473    }
474
475    fn dismiss_by_suggestion(&self, suggestion: &Suggestion) -> Result<()> {
476        if let Some(key) = suggestion.dismissal_key() {
477            match suggestion {
478                Suggestion::Dynamic {
479                    suggestion_type, ..
480                } => self
481                    .dbs()?
482                    .writer
483                    .write(|dao| dao.insert_dynamic_dismissal(suggestion_type, key))?,
484                _ => self.dismiss_by_key(key)?,
485            }
486        }
487        Ok(())
488    }
489
490    fn dismiss_by_key(&self, key: &str) -> Result<()> {
491        self.dbs()?.writer.write(|dao| dao.insert_dismissal(key))
492    }
493
494    fn dismiss_suggestion(&self, suggestion_url: String) -> Result<()> {
495        self.dbs()?
496            .writer
497            .write(|dao| dao.insert_dismissal(&suggestion_url))
498    }
499
500    fn clear_dismissed_suggestions(&self) -> Result<()> {
501        self.dbs()?.writer.write(|dao| dao.clear_dismissals())?;
502        Ok(())
503    }
504
505    fn is_dismissed_by_suggestion(&self, suggestion: &Suggestion) -> Result<bool> {
506        if let Some(key) = suggestion.dismissal_key() {
507            match suggestion {
508                Suggestion::Dynamic {
509                    suggestion_type, ..
510                } => self
511                    .dbs()?
512                    .reader
513                    .read(|dao| dao.has_dynamic_dismissal(suggestion_type, key)),
514                _ => self.dbs()?.reader.read(|dao| dao.has_dismissal(key)),
515            }
516        } else {
517            Ok(false)
518        }
519    }
520
521    fn is_dismissed_by_key(&self, key: &str) -> Result<bool> {
522        self.dbs()?.reader.read(|dao| dao.has_dismissal(key))
523    }
524
525    fn any_dismissed_suggestions(&self) -> Result<bool> {
526        self.dbs()?.reader.read(|dao| dao.any_dismissals())
527    }
528
529    fn interrupt(&self, kind: Option<InterruptKind>) {
530        if let Some(dbs) = self.dbs.get() {
531            // Only interrupt if the databases are already open.
532            match kind.unwrap_or(InterruptKind::Read) {
533                InterruptKind::Read => {
534                    dbs.reader.interrupt_handle.interrupt();
535                }
536                InterruptKind::Write => {
537                    dbs.writer.interrupt_handle.interrupt();
538                }
539                InterruptKind::ReadWrite => {
540                    dbs.reader.interrupt_handle.interrupt();
541                    dbs.writer.interrupt_handle.interrupt();
542                }
543            }
544        }
545    }
546
547    fn clear(&self) -> Result<()> {
548        self.dbs()?.writer.write(|dao| dao.clear())
549    }
550
551    pub fn fetch_global_config(&self) -> Result<SuggestGlobalConfig> {
552        self.dbs()?.reader.read(|dao| dao.get_global_config())
553    }
554
555    pub fn fetch_provider_config(
556        &self,
557        provider: SuggestionProvider,
558    ) -> Result<Option<SuggestProviderConfig>> {
559        self.dbs()?
560            .reader
561            .read(|dao| dao.get_provider_config(provider))
562    }
563
564    // Cause the next ingestion to re-ingest all data
565    pub fn force_reingest(&self) {
566        let writer = &self.dbs().unwrap().writer;
567        writer.write(|dao| dao.force_reingest()).unwrap();
568    }
569
570    fn fetch_geonames(
571        &self,
572        query: &str,
573        match_name_prefix: bool,
574        filter: Option<Vec<Geoname>>,
575    ) -> Result<Vec<GeonameMatch>> {
576        self.dbs()?.reader.read(|dao| {
577            dao.fetch_geonames(
578                query,
579                match_name_prefix,
580                filter.as_ref().map(|f| f.iter().collect()),
581            )
582        })
583    }
584
585    pub fn fetch_geoname_alternates(&self, geoname: &Geoname) -> Result<GeonameAlternates> {
586        self.dbs()?
587            .reader
588            .read(|dao| dao.fetch_geoname_alternates(geoname))
589    }
590}
591
592impl<S> SuggestStoreInner<S>
593where
594    S: Client,
595{
596    pub fn ingest(
597        &self,
598        constraints: SuggestIngestionConstraints,
599    ) -> Result<SuggestIngestionMetrics> {
600        breadcrumb!("Ingestion starting");
601        let writer = &self.dbs()?.writer;
602        let mut metrics = SuggestIngestionMetrics::default();
603        if constraints.empty_only && !writer.read(|dao| dao.suggestions_table_empty())? {
604            return Ok(metrics);
605        }
606
607        // Figure out which record types we're ingesting and group them by
608        // collection. A record type may be used by multiple providers, but we
609        // want to ingest each one at most once. We always ingest some types
610        // like global config.
611        let mut record_types_by_collection = HashMap::from([(
612            Collection::Other,
613            HashSet::from([SuggestRecordType::GlobalConfig]),
614        )]);
615        for provider in constraints
616            .providers
617            .as_ref()
618            .unwrap_or(&DEFAULT_INGEST_PROVIDERS.to_vec())
619            .iter()
620        {
621            for (collection, provider_rts) in provider.record_types_by_collection() {
622                record_types_by_collection
623                    .entry(collection)
624                    .or_default()
625                    .extend(provider_rts.into_iter());
626            }
627        }
628
629        // Create a single write scope for all DB operations
630        let mut write_scope = writer.write_scope()?;
631
632        // Read the previously ingested records.  We use this to calculate what's changed
633        let ingested_records = write_scope.read(|dao| dao.get_ingested_records())?;
634
635        // For each collection, fetch all records
636        for (collection, record_types) in record_types_by_collection {
637            breadcrumb!("Ingesting collection {}", collection.name());
638            let records = self.settings_client.get_records(collection)?;
639
640            // For each record type in that collection, calculate the changes and pass them to
641            // [Self::ingest_records]
642            for record_type in record_types {
643                breadcrumb!("Ingesting record_type: {record_type}");
644                metrics.measure_ingest(record_type.to_string(), |context| {
645                    let changes = RecordChanges::new(
646                        records.iter().filter(|r| r.record_type() == record_type),
647                        ingested_records.iter().filter(|i| {
648                            i.record_type == record_type.as_str()
649                                && i.collection == collection.name()
650                        }),
651                    );
652                    write_scope.write(|dao| {
653                        self.process_changes(dao, collection, changes, &constraints, context)
654                    })
655                })?;
656                write_scope.err_if_interrupted()?;
657            }
658        }
659        breadcrumb!("Ingestion complete");
660
661        Ok(metrics)
662    }
663
664    fn process_changes(
665        &self,
666        dao: &mut SuggestDao,
667        collection: Collection,
668        changes: RecordChanges<'_>,
669        constraints: &SuggestIngestionConstraints,
670        context: &mut MetricsContext,
671    ) -> Result<()> {
672        for record in &changes.new {
673            trace!("Ingesting record ID: {}", record.id.as_str());
674            self.process_record(dao, record, constraints, context)?;
675        }
676        for record in &changes.updated {
677            // Drop any data that we previously ingested from this record.
678            // Suggestions in particular don't have a stable identifier, and
679            // determining which suggestions in the record actually changed is
680            // more complicated than dropping and re-ingesting all of them.
681            trace!("Reingesting updated record ID: {}", record.id.as_str());
682            dao.delete_record_data(&record.id)?;
683            self.process_record(dao, record, constraints, context)?;
684        }
685        for record in &changes.unchanged {
686            if self.should_reprocess_record(dao, record, constraints)? {
687                trace!("Reingesting unchanged record ID: {}", record.id.as_str());
688                self.process_record(dao, record, constraints, context)?;
689            } else {
690                trace!("Skipping unchanged record ID: {}", record.id.as_str());
691            }
692        }
693        for record in &changes.deleted {
694            trace!("Deleting record ID: {:?}", record.id);
695            dao.delete_record_data(&record.id)?;
696        }
697        dao.update_ingested_records(
698            collection.name(),
699            &changes.new,
700            &changes.updated,
701            &changes.deleted,
702        )?;
703        Ok(())
704    }
705
706    fn process_record(
707        &self,
708        dao: &mut SuggestDao,
709        record: &Record,
710        constraints: &SuggestIngestionConstraints,
711        context: &mut MetricsContext,
712    ) -> Result<()> {
713        match &record.payload {
714            SuggestRecord::Amp => {
715                self.download_attachment(dao, record, context, |dao, record_id, suggestions| {
716                    dao.insert_amp_suggestions(
717                        record_id,
718                        suggestions,
719                        constraints.amp_matching_uses_fts(),
720                    )
721                })?;
722            }
723            SuggestRecord::Wikipedia => {
724                self.download_attachment(dao, record, context, |dao, record_id, suggestions| {
725                    dao.insert_wikipedia_suggestions(record_id, suggestions)
726                })?;
727            }
728            SuggestRecord::Icon => {
729                let (Some(icon_id), Some(attachment)) =
730                    (record.id.as_icon_id(), record.attachment.as_ref())
731                else {
732                    // An icon record should have an icon ID and an
733                    // attachment. Icons that don't have these are
734                    // malformed, so skip to the next record.
735                    return Ok(());
736                };
737                let data = context
738                    .measure_download(|| self.settings_client.download_attachment(record))?;
739                dao.put_icon(icon_id, &data, &attachment.mimetype)?;
740            }
741            SuggestRecord::Amo => {
742                self.download_attachment(dao, record, context, |dao, record_id, suggestions| {
743                    dao.insert_amo_suggestions(record_id, suggestions)
744                })?;
745            }
746            SuggestRecord::Yelp => {
747                self.download_attachment(dao, record, context, |dao, record_id, suggestions| {
748                    match suggestions.first() {
749                        Some(suggestion) => dao.insert_yelp_suggestions(record_id, suggestion),
750                        None => Ok(()),
751                    }
752                })?;
753            }
754            SuggestRecord::Mdn => {
755                self.download_attachment(dao, record, context, |dao, record_id, suggestions| {
756                    dao.insert_mdn_suggestions(record_id, suggestions)
757                })?;
758            }
759            SuggestRecord::Weather => self.process_weather_record(dao, record, context)?,
760            SuggestRecord::GlobalConfig(config) => {
761                dao.put_global_config(&SuggestGlobalConfig::from(config))?
762            }
763            SuggestRecord::Fakespot => {
764                self.download_attachment(dao, record, context, |dao, record_id, suggestions| {
765                    dao.insert_fakespot_suggestions(record_id, suggestions)
766                })?;
767            }
768            SuggestRecord::Dynamic(r) => {
769                if constraints.matches_dynamic_record(r) {
770                    self.download_attachment(
771                        dao,
772                        record,
773                        context,
774                        |dao, record_id, suggestions| {
775                            dao.insert_dynamic_suggestions(record_id, r, suggestions)
776                        },
777                    )?;
778                }
779            }
780            SuggestRecord::Geonames => self.process_geonames_record(dao, record, context)?,
781            SuggestRecord::GeonamesAlternates => {
782                self.process_geonames_alternates_record(dao, record, context)?
783            }
784        }
785        Ok(())
786    }
787
788    pub(crate) fn download_attachment<T>(
789        &self,
790        dao: &mut SuggestDao,
791        record: &Record,
792        context: &mut MetricsContext,
793        ingestion_handler: impl FnOnce(&mut SuggestDao<'_>, &SuggestRecordId, &[T]) -> Result<()>,
794    ) -> Result<()>
795    where
796        T: DeserializeOwned,
797    {
798        if record.attachment.is_none() {
799            return Ok(());
800        };
801
802        let attachment_data =
803            context.measure_download(|| self.settings_client.download_attachment(record))?;
804        match serde_json::from_slice::<SuggestAttachment<T>>(&attachment_data) {
805            Ok(attachment) => ingestion_handler(dao, &record.id, attachment.suggestions()),
806            // If the attachment doesn't match our expected schema, just skip it.  It's possible
807            // that we're using an older version.  If so, we'll get the data when we re-ingest
808            // after updating the schema.
809            Err(_) => Ok(()),
810        }
811    }
812
813    fn should_reprocess_record(
814        &self,
815        dao: &mut SuggestDao,
816        record: &Record,
817        constraints: &SuggestIngestionConstraints,
818    ) -> Result<bool> {
819        match &record.payload {
820            SuggestRecord::Dynamic(r) => Ok(!dao
821                .are_suggestions_ingested_for_record(&record.id)?
822                && constraints.matches_dynamic_record(r)),
823            SuggestRecord::Amp => {
824                Ok(constraints.amp_matching_uses_fts()
825                    && !dao.is_amp_fts_data_ingested(&record.id)?)
826            }
827            _ => Ok(false),
828        }
829    }
830}
831
832/// Tracks changes in suggest records since the last ingestion
833struct RecordChanges<'a> {
834    new: Vec<&'a Record>,
835    updated: Vec<&'a Record>,
836    deleted: Vec<&'a IngestedRecord>,
837    unchanged: Vec<&'a Record>,
838}
839
840impl<'a> RecordChanges<'a> {
841    fn new(
842        current: impl Iterator<Item = &'a Record>,
843        previously_ingested: impl Iterator<Item = &'a IngestedRecord>,
844    ) -> Self {
845        let mut ingested_map: HashMap<&str, &IngestedRecord> =
846            previously_ingested.map(|i| (i.id.as_str(), i)).collect();
847        // Iterate through current, finding new/updated records.
848        // Remove existing records from ingested_map.
849        let mut new = vec![];
850        let mut updated = vec![];
851        let mut unchanged = vec![];
852        for r in current {
853            match ingested_map.entry(r.id.as_str()) {
854                Entry::Vacant(_) => new.push(r),
855                Entry::Occupied(e) => {
856                    if e.remove().last_modified != r.last_modified {
857                        updated.push(r);
858                    } else {
859                        unchanged.push(r);
860                    }
861                }
862            }
863        }
864        // Anything left in ingested_map is a deleted record
865        let deleted = ingested_map.into_values().collect();
866        Self {
867            new,
868            deleted,
869            updated,
870            unchanged,
871        }
872    }
873}
874
875#[cfg(feature = "benchmark_api")]
876impl<S> SuggestStoreInner<S>
877where
878    S: Client,
879{
880    pub fn into_settings_client(self) -> S {
881        self.settings_client
882    }
883
884    pub fn ensure_db_initialized(&self) {
885        self.dbs().unwrap();
886    }
887
888    fn checkpoint(&self) {
889        let conn = self.dbs().unwrap().writer.conn.lock();
890        conn.pragma_update(None, "wal_checkpoint", "TRUNCATE")
891            .expect("Error performing checkpoint");
892    }
893
894    pub fn ingest_records_by_type(
895        &self,
896        collection: Collection,
897        ingest_record_type: SuggestRecordType,
898    ) {
899        let writer = &self.dbs().unwrap().writer;
900        let mut context = MetricsContext::default();
901        let ingested_records = writer.read(|dao| dao.get_ingested_records()).unwrap();
902        let records = self.settings_client.get_records(collection).unwrap();
903
904        let changes = RecordChanges::new(
905            records
906                .iter()
907                .filter(|r| r.record_type() == ingest_record_type),
908            ingested_records
909                .iter()
910                .filter(|i| i.record_type == ingest_record_type.as_str()),
911        );
912        writer
913            .write(|dao| {
914                self.process_changes(
915                    dao,
916                    collection,
917                    changes,
918                    &SuggestIngestionConstraints::default(),
919                    &mut context,
920                )
921            })
922            .unwrap();
923    }
924
925    pub fn table_row_counts(&self) -> Vec<(String, u32)> {
926        use sql_support::ConnExt;
927
928        // Note: since this is just used for debugging, use unwrap to simplify the error handling.
929        let reader = &self.dbs().unwrap().reader;
930        let conn = reader.conn.lock();
931        let table_names: Vec<String> = conn
932            .query_rows_and_then(
933                "SELECT name FROM sqlite_master where type = 'table'",
934                (),
935                |row| row.get(0),
936            )
937            .unwrap();
938        let mut table_names_with_counts: Vec<(String, u32)> = table_names
939            .into_iter()
940            .map(|name| {
941                let count: u32 = conn
942                    .conn_ext_query_one(&format!("SELECT COUNT(*) FROM {name}"))
943                    .unwrap();
944                (name, count)
945            })
946            .collect();
947        table_names_with_counts.sort_by(|a, b| (b.1.cmp(&a.1)));
948        table_names_with_counts
949    }
950
951    pub fn db_size(&self) -> usize {
952        use sql_support::ConnExt;
953
954        let reader = &self.dbs().unwrap().reader;
955        let conn = reader.conn.lock();
956        conn.conn_ext_query_one(
957            "SELECT page_size * page_count FROM pragma_page_count(), pragma_page_size()",
958        )
959        .unwrap()
960    }
961}
962
963/// Holds a store's open connections to the Suggest database.
964struct SuggestStoreDbs {
965    /// A read-write connection used to update the database with new data.
966    writer: SuggestDb,
967    /// A read-only connection used to query the database.
968    reader: SuggestDb,
969}
970
971impl SuggestStoreDbs {
972    fn open(path: &Path, extensions_to_load: &[Sqlite3Extension]) -> Result<Self> {
973        // Order is important here: the writer must be opened first, so that it
974        // can set up the database and run any migrations.
975        let writer = SuggestDb::open(path, extensions_to_load, ConnectionType::ReadWrite)?;
976        let reader = SuggestDb::open(path, extensions_to_load, ConnectionType::ReadOnly)?;
977        Ok(Self { writer, reader })
978    }
979}
980
981#[cfg(test)]
982pub(crate) mod tests {
983    use super::*;
984    use crate::suggestion::YelpSubjectType;
985
986    use std::sync::atomic::{AtomicUsize, Ordering};
987
988    use crate::{
989        db::DEFAULT_SUGGESTION_SCORE, provider::AmpMatchingStrategy, suggestion::FtsMatchInfo,
990        testing::*, SuggestionProvider,
991    };
992
993    // Extra methods for the tests
994    impl SuggestIngestionConstraints {
995        fn amp_with_fts() -> Self {
996            Self {
997                providers: Some(vec![SuggestionProvider::Amp]),
998                provider_constraints: Some(SuggestionProviderConstraints {
999                    amp_alternative_matching: Some(AmpMatchingStrategy::FtsAgainstFullKeywords),
1000                    ..SuggestionProviderConstraints::default()
1001                }),
1002                ..Self::default()
1003            }
1004        }
1005        fn amp_without_fts() -> Self {
1006            Self {
1007                providers: Some(vec![SuggestionProvider::Amp]),
1008                ..Self::default()
1009            }
1010        }
1011    }
1012
1013    /// In-memory Suggest store for testing
1014    pub(crate) struct TestStore {
1015        pub inner: SuggestStoreInner<MockRemoteSettingsClient>,
1016    }
1017
1018    impl TestStore {
1019        pub fn new(client: MockRemoteSettingsClient) -> Self {
1020            static COUNTER: AtomicUsize = AtomicUsize::new(0);
1021            let db_path = format!(
1022                "file:test_store_data_{}?mode=memory&cache=shared",
1023                COUNTER.fetch_add(1, Ordering::Relaxed),
1024            );
1025            Self {
1026                inner: SuggestStoreInner::new(db_path, vec![], client),
1027            }
1028        }
1029
1030        pub fn client_mut(&mut self) -> &mut MockRemoteSettingsClient {
1031            &mut self.inner.settings_client
1032        }
1033
1034        pub fn read<T>(&self, op: impl FnOnce(&SuggestDao) -> Result<T>) -> Result<T> {
1035            self.inner.dbs().unwrap().reader.read(op)
1036        }
1037
1038        pub fn write<T>(&self, op: impl FnMut(&mut SuggestDao) -> Result<T>) -> Result<T> {
1039            self.inner.dbs().unwrap().writer.write(op)
1040        }
1041
1042        pub fn count_rows(&self, table_name: &str) -> u64 {
1043            let sql = format!("SELECT count(*) FROM {table_name}");
1044            self.read(|dao| Ok(dao.conn.conn_ext_query_one(&sql)?))
1045                .unwrap_or_else(|e| panic!("SQL error in count: {e}"))
1046        }
1047
1048        pub fn ingest(&self, constraints: SuggestIngestionConstraints) {
1049            self.inner.ingest(constraints).unwrap();
1050        }
1051
1052        pub fn fetch_suggestions(&self, query: SuggestionQuery) -> Vec<Suggestion> {
1053            self.inner.query(query).unwrap().suggestions
1054        }
1055
1056        pub fn fetch_global_config(&self) -> SuggestGlobalConfig {
1057            self.inner
1058                .fetch_global_config()
1059                .expect("Error fetching global config")
1060        }
1061
1062        pub fn fetch_provider_config(
1063            &self,
1064            provider: SuggestionProvider,
1065        ) -> Option<SuggestProviderConfig> {
1066            self.inner
1067                .fetch_provider_config(provider)
1068                .expect("Error fetching provider config")
1069        }
1070
1071        pub fn fetch_geonames(
1072            &self,
1073            query: &str,
1074            match_name_prefix: bool,
1075            filter: Option<Vec<Geoname>>,
1076        ) -> Vec<GeonameMatch> {
1077            self.inner
1078                .fetch_geonames(query, match_name_prefix, filter)
1079                .expect("Error fetching geonames")
1080        }
1081    }
1082
1083    /// Tests that `SuggestStore` is usable with UniFFI, which requires exposed
1084    /// interfaces to be `Send` and `Sync`.
1085    #[test]
1086    fn is_thread_safe() {
1087        before_each();
1088
1089        fn is_send_sync<T: Send + Sync>() {}
1090        is_send_sync::<SuggestStore>();
1091    }
1092
1093    /// Tests ingesting suggestions into an empty database.
1094    #[test]
1095    fn ingest_suggestions() -> anyhow::Result<()> {
1096        before_each();
1097
1098        let store = TestStore::new(
1099            MockRemoteSettingsClient::default()
1100                .with_record(SuggestionProvider::Amp.record("1234", json![los_pollos_amp()]))
1101                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
1102        );
1103        store.ingest(SuggestIngestionConstraints::all_providers());
1104        assert_eq!(
1105            store.fetch_suggestions(SuggestionQuery::amp("lo")),
1106            vec![los_pollos_suggestion("los pollos", None)],
1107        );
1108        Ok(())
1109    }
1110
1111    /// Tests ingesting suggestions into an empty database.
1112    #[test]
1113    fn ingest_empty_only() -> anyhow::Result<()> {
1114        before_each();
1115
1116        let mut store = TestStore::new(
1117            MockRemoteSettingsClient::default()
1118                .with_record(SuggestionProvider::Amp.record("1234", json![los_pollos_amp()])),
1119        );
1120        // suggestions_table_empty returns true before the ingestion is complete
1121        assert!(store.read(|dao| dao.suggestions_table_empty())?);
1122        // This ingestion should run, since the DB is empty
1123        store.ingest(SuggestIngestionConstraints {
1124            empty_only: true,
1125            ..SuggestIngestionConstraints::all_providers()
1126        });
1127        // suggestions_table_empty returns false after the ingestion is complete
1128        assert!(!store.read(|dao| dao.suggestions_table_empty())?);
1129
1130        // This ingestion should not run since the DB is no longer empty
1131        store.client_mut().update_record(
1132            SuggestionProvider::Amp
1133                .record("1234", json!([los_pollos_amp(), good_place_eats_amp()])),
1134        );
1135
1136        store.ingest(SuggestIngestionConstraints {
1137            empty_only: true,
1138            ..SuggestIngestionConstraints::all_providers()
1139        });
1140        // "la" should not match the good place eats suggestion, since that should not have been
1141        // ingested.
1142        assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("la")), vec![]);
1143
1144        Ok(())
1145    }
1146
1147    /// Tests ingesting suggestions with icons.
1148    #[test]
1149    fn ingest_amp_icons() -> anyhow::Result<()> {
1150        before_each();
1151
1152        let store = TestStore::new(
1153            MockRemoteSettingsClient::default()
1154                .with_record(
1155                    SuggestionProvider::Amp
1156                        .record("1234", json!([los_pollos_amp(), good_place_eats_amp()])),
1157                )
1158                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
1159                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon())),
1160        );
1161        // This ingestion should run, since the DB is empty
1162        store.ingest(SuggestIngestionConstraints::all_providers());
1163
1164        assert_eq!(
1165            store.fetch_suggestions(SuggestionQuery::amp("lo")),
1166            vec![los_pollos_suggestion("los pollos", None)]
1167        );
1168        assert_eq!(
1169            store.fetch_suggestions(SuggestionQuery::amp("la")),
1170            vec![good_place_eats_suggestion("lasagna", None)]
1171        );
1172
1173        Ok(())
1174    }
1175
1176    #[test]
1177    fn ingest_amp_full_keywords() -> anyhow::Result<()> {
1178        before_each();
1179
1180        let store = TestStore::new(MockRemoteSettingsClient::default()
1181            .with_record(
1182                SuggestionProvider::Amp.record("1234", json!([
1183                // AMP attachment with full keyword data
1184                los_pollos_amp().merge(json!({
1185                    "keywords": ["lo", "los", "los p", "los pollos", "los pollos h", "los pollos hermanos"],
1186                    "full_keywords": [
1187                        // Full keyword for the first 4 keywords
1188                        ("los pollos", 4),
1189                        // Full keyword for the next 2 keywords
1190                        ("los pollos hermanos (restaurant)", 2),
1191                    ],
1192                })),
1193                // AMP attachment without full keyword data
1194                good_place_eats_amp().remove("full_keywords"),
1195            ])))
1196            .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
1197            .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon()))
1198        );
1199        store.ingest(SuggestIngestionConstraints::all_providers());
1200
1201        // (query string, expected suggestion, expected dismissal key)
1202        let tests = [
1203            (
1204                "lo",
1205                los_pollos_suggestion("los pollos", None),
1206                Some("los pollos"),
1207            ),
1208            (
1209                "los pollos",
1210                los_pollos_suggestion("los pollos", None),
1211                Some("los pollos"),
1212            ),
1213            (
1214                "los pollos h",
1215                los_pollos_suggestion("los pollos hermanos (restaurant)", None),
1216                Some("los pollos hermanos (restaurant)"),
1217            ),
1218            (
1219                "la",
1220                good_place_eats_suggestion("", None),
1221                Some("https://www.lasagna.restaurant"),
1222            ),
1223            (
1224                "lasagna",
1225                good_place_eats_suggestion("", None),
1226                Some("https://www.lasagna.restaurant"),
1227            ),
1228            (
1229                "lasagna come out tomorrow",
1230                good_place_eats_suggestion("", None),
1231                Some("https://www.lasagna.restaurant"),
1232            ),
1233        ];
1234        for (query, expected_suggestion, expected_dismissal_key) in tests {
1235            // Do a query and check the returned suggestions.
1236            let suggestions = store.fetch_suggestions(SuggestionQuery::amp(query));
1237            assert_eq!(suggestions, vec![expected_suggestion.clone()]);
1238
1239            // Check the returned suggestion's dismissal key.
1240            assert_eq!(suggestions[0].dismissal_key(), expected_dismissal_key);
1241
1242            // Dismiss the suggestion.
1243            let dismissal_key = suggestions[0].dismissal_key().unwrap();
1244            store.inner.dismiss_by_suggestion(&suggestions[0])?;
1245            assert_eq!(store.fetch_suggestions(SuggestionQuery::amp(query)), vec![]);
1246            assert!(store.inner.is_dismissed_by_suggestion(&suggestions[0])?);
1247            assert!(store.inner.is_dismissed_by_key(dismissal_key)?);
1248            assert!(store.inner.any_dismissed_suggestions()?);
1249
1250            // Clear dismissals and fetch again.
1251            store.inner.clear_dismissed_suggestions()?;
1252            assert_eq!(
1253                store.fetch_suggestions(SuggestionQuery::amp(query)),
1254                vec![expected_suggestion.clone()]
1255            );
1256            assert!(!store.inner.is_dismissed_by_suggestion(&suggestions[0])?);
1257            assert!(!store.inner.is_dismissed_by_key(dismissal_key)?);
1258            assert!(!store.inner.any_dismissed_suggestions()?);
1259
1260            // Dismiss the suggestion by its dismissal key.
1261            store.inner.dismiss_by_key(dismissal_key)?;
1262            assert_eq!(store.fetch_suggestions(SuggestionQuery::amp(query)), vec![]);
1263            assert!(store.inner.is_dismissed_by_suggestion(&suggestions[0])?);
1264            assert!(store.inner.is_dismissed_by_key(dismissal_key)?);
1265            assert!(store.inner.any_dismissed_suggestions()?);
1266
1267            // Clear dismissals and fetch again.
1268            store.inner.clear_dismissed_suggestions()?;
1269            assert_eq!(
1270                store.fetch_suggestions(SuggestionQuery::amp(query)),
1271                vec![expected_suggestion.clone()]
1272            );
1273            assert!(!store.inner.is_dismissed_by_suggestion(&suggestions[0])?);
1274            assert!(!store.inner.is_dismissed_by_key(dismissal_key)?);
1275            assert!(!store.inner.any_dismissed_suggestions()?);
1276
1277            // Dismiss the suggestion by its raw URL using the deprecated API.
1278            let raw_url = expected_suggestion.raw_url().unwrap();
1279            store.inner.dismiss_suggestion(raw_url.to_string())?;
1280            assert_eq!(store.fetch_suggestions(SuggestionQuery::amp(query)), vec![]);
1281            assert!(store.inner.is_dismissed_by_key(raw_url)?);
1282            assert!(store.inner.any_dismissed_suggestions()?);
1283
1284            // Clear dismissals and fetch again.
1285            store.inner.clear_dismissed_suggestions()?;
1286            assert_eq!(
1287                store.fetch_suggestions(SuggestionQuery::amp(query)),
1288                vec![expected_suggestion.clone()]
1289            );
1290            assert!(!store.inner.is_dismissed_by_suggestion(&suggestions[0])?);
1291            assert!(!store.inner.is_dismissed_by_key(dismissal_key)?);
1292            assert!(!store.inner.is_dismissed_by_key(raw_url)?);
1293            assert!(!store.inner.any_dismissed_suggestions()?);
1294        }
1295
1296        Ok(())
1297    }
1298
1299    #[test]
1300    fn ingest_wikipedia_full_keywords() -> anyhow::Result<()> {
1301        before_each();
1302
1303        let store = TestStore::new(
1304            MockRemoteSettingsClient::default()
1305                .with_record(SuggestionProvider::Wikipedia.record(
1306                    "1234",
1307                    json!([
1308                        // Wikipedia attachment with full keyword data.  We should ignore the full
1309                        // keyword data for Wikipedia suggestions
1310                        california_wiki(),
1311                        // california_wiki().merge(json!({
1312                        //     "keywords": ["cal", "cali", "california"],
1313                        //     "full_keywords": [("california institute of technology", 3)],
1314                        // })),
1315                    ]),
1316                ))
1317                .with_record(SuggestionProvider::Wikipedia.icon(california_icon())),
1318        );
1319        store.ingest(SuggestIngestionConstraints::all_providers());
1320
1321        assert_eq!(
1322            store.fetch_suggestions(SuggestionQuery::wikipedia("cal")),
1323            // Even though this had a full_keywords field, we should ignore it since it's a
1324            // wikipedia suggestion and use the keywords.rs code instead
1325            vec![california_suggestion("california")],
1326        );
1327
1328        Ok(())
1329    }
1330
1331    #[test]
1332    fn amp_no_keyword_expansion() -> anyhow::Result<()> {
1333        before_each();
1334
1335        let store = TestStore::new(
1336            MockRemoteSettingsClient::default()
1337                // Setup the keywords such that:
1338                //   * There's a `chicken` keyword, which is not a substring of any full
1339                //     keywords (i.e. it was the result of keyword expansion).
1340                //   * There's a `los pollos ` keyword with an extra space
1341                .with_record(
1342                    SuggestionProvider::Amp.record(
1343                    "1234",
1344                    los_pollos_amp().merge(json!({
1345                        "keywords": ["los", "los pollos", "los pollos ", "los pollos hermanos", "chicken"],
1346                        "full_keywords": [("los pollos", 3), ("los pollos hermanos", 2)],
1347                    }))
1348                ))
1349                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
1350        );
1351        store.ingest(SuggestIngestionConstraints::all_providers());
1352        assert_eq!(
1353            store.fetch_suggestions(SuggestionQuery {
1354                provider_constraints: Some(SuggestionProviderConstraints {
1355                    amp_alternative_matching: Some(AmpMatchingStrategy::NoKeywordExpansion),
1356                    ..SuggestionProviderConstraints::default()
1357                }),
1358                // Should not match, because `chicken` is not a substring of a full keyword.
1359                // i.e. it was added because of keyword expansion.
1360                ..SuggestionQuery::amp("chicken")
1361            }),
1362            vec![],
1363        );
1364        assert_eq!(
1365            store.fetch_suggestions(SuggestionQuery {
1366                provider_constraints: Some(SuggestionProviderConstraints {
1367                    amp_alternative_matching: Some(AmpMatchingStrategy::NoKeywordExpansion),
1368                    ..SuggestionProviderConstraints::default()
1369                }),
1370                // Should match, even though "los pollos " technically is not a substring
1371                // because there's an extra space.  The reason these keywords are in the DB is
1372                // because we want to keep showing the current suggestion when the user types
1373                // the space key.
1374                ..SuggestionQuery::amp("los pollos ")
1375            }),
1376            vec![los_pollos_suggestion("los pollos", None)],
1377        );
1378        Ok(())
1379    }
1380
1381    #[test]
1382    fn amp_fts_against_full_keywords() -> anyhow::Result<()> {
1383        before_each();
1384
1385        let store = TestStore::new(
1386            MockRemoteSettingsClient::default()
1387                // Make sure there's full keywords to match against
1388                .with_record(SuggestionProvider::Amp.record(
1389                    "1234",
1390                    los_pollos_amp().merge(json!({
1391                        "keywords": ["los", "los pollos", "los pollos ", "los pollos hermanos"],
1392                        "full_keywords": [("los pollos", 3), ("los pollos hermanos", 1)],
1393                    })),
1394                ))
1395                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
1396        );
1397        store.ingest(SuggestIngestionConstraints::amp_with_fts());
1398        assert_eq!(
1399            store.fetch_suggestions(SuggestionQuery {
1400                provider_constraints: Some(SuggestionProviderConstraints {
1401                    amp_alternative_matching: Some(AmpMatchingStrategy::FtsAgainstFullKeywords),
1402                    ..SuggestionProviderConstraints::default()
1403                }),
1404                // "Hermanos" should match, even though it's not listed in the keywords,
1405                // because this strategy uses an FTS match against the full keyword list.
1406                ..SuggestionQuery::amp("hermanos")
1407            }),
1408            vec![los_pollos_suggestion(
1409                "hermanos",
1410                Some(FtsMatchInfo {
1411                    prefix: false,
1412                    stemming: false,
1413                })
1414            )],
1415        );
1416        Ok(())
1417    }
1418
1419    #[test]
1420    fn amp_fts_against_title() -> anyhow::Result<()> {
1421        before_each();
1422
1423        let store = TestStore::new(
1424            MockRemoteSettingsClient::default()
1425                .with_record(SuggestionProvider::Amp.record("1234", los_pollos_amp()))
1426                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
1427        );
1428        store.ingest(SuggestIngestionConstraints::amp_with_fts());
1429        assert_eq!(
1430            store.fetch_suggestions(SuggestionQuery {
1431                provider_constraints: Some(SuggestionProviderConstraints {
1432                    amp_alternative_matching: Some(AmpMatchingStrategy::FtsAgainstTitle),
1433                    ..SuggestionProviderConstraints::default()
1434                }),
1435                // "Albuquerque" should match, even though it's not listed in the keywords,
1436                // because this strategy uses an FTS match against the title
1437                ..SuggestionQuery::amp("albuquerque")
1438            }),
1439            vec![los_pollos_suggestion(
1440                "albuquerque",
1441                Some(FtsMatchInfo {
1442                    prefix: false,
1443                    stemming: false,
1444                })
1445            )],
1446        );
1447        Ok(())
1448    }
1449
1450    /// Tests ingesting a data attachment containing a single suggestion,
1451    /// instead of an array of suggestions.
1452    #[test]
1453    fn ingest_one_suggestion_in_data_attachment() -> anyhow::Result<()> {
1454        before_each();
1455
1456        let store = TestStore::new(
1457            MockRemoteSettingsClient::default()
1458                // This record contains just one JSON object, rather than an array of them
1459                .with_record(SuggestionProvider::Amp.record("1234", los_pollos_amp()))
1460                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
1461        );
1462        store.ingest(SuggestIngestionConstraints::all_providers());
1463        assert_eq!(
1464            store.fetch_suggestions(SuggestionQuery::amp("lo")),
1465            vec![los_pollos_suggestion("los pollos", None)],
1466        );
1467
1468        Ok(())
1469    }
1470
1471    /// Tests re-ingesting suggestions from an updated attachment.
1472    #[test]
1473    fn reingest_amp_suggestions() -> anyhow::Result<()> {
1474        before_each();
1475
1476        let mut store = TestStore::new(
1477            MockRemoteSettingsClient::default().with_record(
1478                SuggestionProvider::Amp
1479                    .record("1234", json!([los_pollos_amp(), good_place_eats_amp()])),
1480            ),
1481        );
1482        // Ingest once
1483        store.ingest(SuggestIngestionConstraints::all_providers());
1484        // Update the snapshot with new suggestions: Los pollos has a new name and Good place eats
1485        // is now serving Penne
1486        store
1487            .client_mut()
1488            .update_record(SuggestionProvider::Amp.record(
1489                "1234",
1490                json!([
1491                    los_pollos_amp().merge(json!({
1492                        "title": "Los Pollos Hermanos - Now Serving at 14 Locations!",
1493                    })),
1494                    good_place_eats_amp().merge(json!({
1495                        "keywords": ["pe", "pen", "penne", "penne for your thoughts"],
1496                        "title": "Penne for Your Thoughts",
1497                        "url": "https://penne.biz",
1498                    }))
1499                ]),
1500            ));
1501        store.ingest(SuggestIngestionConstraints::all_providers());
1502
1503        assert!(matches!(
1504            store.fetch_suggestions(SuggestionQuery::amp("lo")).as_slice(),
1505            [ Suggestion::Amp { title, .. } ] if title == "Los Pollos Hermanos - Now Serving at 14 Locations!",
1506        ));
1507
1508        assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("la")), vec![]);
1509        assert!(matches!(
1510            store.fetch_suggestions(SuggestionQuery::amp("pe")).as_slice(),
1511            [ Suggestion::Amp { title, url, .. } ] if title == "Penne for Your Thoughts" && url == "https://penne.biz"
1512        ));
1513
1514        Ok(())
1515    }
1516
1517    #[test]
1518    fn reingest_amp_after_fts_constraint_changes() -> anyhow::Result<()> {
1519        before_each();
1520
1521        // Ingest with FTS enabled, this will populate the FTS table
1522        let store = TestStore::new(
1523            MockRemoteSettingsClient::default()
1524                .with_record(SuggestionProvider::Amp.record(
1525                    "data-1",
1526                    json!([los_pollos_amp().merge(json!({
1527                        "keywords": ["los", "los pollos", "los pollos ", "los pollos hermanos"],
1528                        "full_keywords": [("los pollos", 3), ("los pollos hermanos", 1)],
1529                    }))]),
1530                ))
1531                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
1532        );
1533        // Ingest without FTS
1534        store.ingest(SuggestIngestionConstraints::amp_without_fts());
1535        // Ingest again with FTS
1536        store.ingest(SuggestIngestionConstraints::amp_with_fts());
1537
1538        assert_eq!(
1539            store.fetch_suggestions(SuggestionQuery {
1540                provider_constraints: Some(SuggestionProviderConstraints {
1541                    amp_alternative_matching: Some(AmpMatchingStrategy::FtsAgainstFullKeywords),
1542                    ..SuggestionProviderConstraints::default()
1543                }),
1544                // "Hermanos" should match, even though it's not listed in the keywords,
1545                // because this strategy uses an FTS match against the full keyword list.
1546                ..SuggestionQuery::amp("hermanos")
1547            }),
1548            vec![los_pollos_suggestion(
1549                "hermanos",
1550                Some(FtsMatchInfo {
1551                    prefix: false,
1552                    stemming: false,
1553                }),
1554            )],
1555        );
1556        Ok(())
1557    }
1558
1559    /// Tests re-ingesting icons from an updated attachment.
1560    #[test]
1561    fn reingest_icons() -> anyhow::Result<()> {
1562        before_each();
1563
1564        let mut store = TestStore::new(
1565            MockRemoteSettingsClient::default()
1566                .with_record(
1567                    SuggestionProvider::Amp
1568                        .record("1234", json!([los_pollos_amp(), good_place_eats_amp()])),
1569                )
1570                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
1571                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon())),
1572        );
1573        // This ingestion should run, since the DB is empty
1574        store.ingest(SuggestIngestionConstraints::all_providers());
1575
1576        // Reingest with updated icon data
1577        //  - Los pollos gets new data and a new id
1578        //  - Good place eats gets new data only
1579        store
1580            .client_mut()
1581            .update_record(SuggestionProvider::Amp.record(
1582                "1234",
1583                json!([
1584                    los_pollos_amp().merge(json!({"icon": "1000"})),
1585                    good_place_eats_amp()
1586                ]),
1587            ))
1588            .delete_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
1589            .add_record(SuggestionProvider::Amp.icon(MockIcon {
1590                id: "1000",
1591                data: "new-los-pollos-icon",
1592                ..los_pollos_icon()
1593            }))
1594            .update_record(SuggestionProvider::Amp.icon(MockIcon {
1595                data: "new-good-place-eats-icon",
1596                ..good_place_eats_icon()
1597            }));
1598        store.ingest(SuggestIngestionConstraints::all_providers());
1599
1600        assert!(matches!(
1601            store.fetch_suggestions(SuggestionQuery::amp("lo")).as_slice(),
1602            [ Suggestion::Amp { icon, .. } ] if *icon == Some("new-los-pollos-icon".as_bytes().to_vec())
1603        ));
1604
1605        assert!(matches!(
1606            store.fetch_suggestions(SuggestionQuery::amp("la")).as_slice(),
1607            [ Suggestion::Amp { icon, .. } ] if *icon == Some("new-good-place-eats-icon".as_bytes().to_vec())
1608        ));
1609
1610        Ok(())
1611    }
1612
1613    /// Tests re-ingesting AMO suggestions from an updated attachment.
1614    #[test]
1615    fn reingest_amo_suggestions() -> anyhow::Result<()> {
1616        before_each();
1617
1618        let mut store = TestStore::new(
1619            MockRemoteSettingsClient::default()
1620                .with_record(SuggestionProvider::Amo.record("data-1", json!([relay_amo()])))
1621                .with_record(
1622                    SuggestionProvider::Amo
1623                        .record("data-2", json!([dark_mode_amo(), foxy_guestures_amo()])),
1624                ),
1625        );
1626
1627        store.ingest(SuggestIngestionConstraints::all_providers());
1628
1629        assert_eq!(
1630            store.fetch_suggestions(SuggestionQuery::amo("masking e")),
1631            vec![relay_suggestion()],
1632        );
1633        assert_eq!(
1634            store.fetch_suggestions(SuggestionQuery::amo("night")),
1635            vec![dark_mode_suggestion()],
1636        );
1637        assert_eq!(
1638            store.fetch_suggestions(SuggestionQuery::amo("grammar")),
1639            vec![foxy_guestures_suggestion()],
1640        );
1641
1642        // Update the snapshot with new suggestions: update the second, drop the
1643        // third, and add the fourth.
1644        store
1645            .client_mut()
1646            .update_record(SuggestionProvider::Amo.record("data-1", json!([relay_amo()])))
1647            .update_record(SuggestionProvider::Amo.record(
1648                "data-2",
1649                json!([
1650                    dark_mode_amo().merge(json!({"title": "Updated second suggestion"})),
1651                    new_tab_override_amo(),
1652                ]),
1653            ));
1654        store.ingest(SuggestIngestionConstraints::all_providers());
1655
1656        assert_eq!(
1657            store.fetch_suggestions(SuggestionQuery::amo("masking e")),
1658            vec![relay_suggestion()],
1659        );
1660        assert!(matches!(
1661            store.fetch_suggestions(SuggestionQuery::amo("night")).as_slice(),
1662            [Suggestion::Amo { title, .. } ] if title == "Updated second suggestion"
1663        ));
1664        assert_eq!(
1665            store.fetch_suggestions(SuggestionQuery::amo("grammar")),
1666            vec![],
1667        );
1668        assert_eq!(
1669            store.fetch_suggestions(SuggestionQuery::amo("image search")),
1670            vec![new_tab_override_suggestion()],
1671        );
1672
1673        Ok(())
1674    }
1675
1676    /// Tests ingestion when previously-ingested suggestions/icons have been deleted.
1677    #[test]
1678    fn ingest_with_deletions() -> anyhow::Result<()> {
1679        before_each();
1680
1681        let mut store = TestStore::new(
1682            MockRemoteSettingsClient::default()
1683                .with_record(SuggestionProvider::Amp.record("data-1", json!([los_pollos_amp()])))
1684                .with_record(
1685                    SuggestionProvider::Amp.record("data-2", json!([good_place_eats_amp()])),
1686                )
1687                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
1688                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon())),
1689        );
1690        store.ingest(SuggestIngestionConstraints::all_providers());
1691        assert_eq!(
1692            store.fetch_suggestions(SuggestionQuery::amp("lo")),
1693            vec![los_pollos_suggestion("los pollos", None)],
1694        );
1695        assert_eq!(
1696            store.fetch_suggestions(SuggestionQuery::amp("la")),
1697            vec![good_place_eats_suggestion("lasagna", None)],
1698        );
1699        // Re-ingest without los-pollos and good place eat's icon.  The suggest store should
1700        // recognize that they're missing and delete them.
1701        store
1702            .client_mut()
1703            .delete_record(SuggestionProvider::Amp.empty_record("data-1"))
1704            .delete_record(SuggestionProvider::Amp.icon(good_place_eats_icon()));
1705        store.ingest(SuggestIngestionConstraints::all_providers());
1706
1707        assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("lo")), vec![]);
1708        assert!(matches!(
1709            store.fetch_suggestions(SuggestionQuery::amp("la")).as_slice(),
1710            [
1711                Suggestion::Amp { icon, icon_mimetype, .. }
1712            ] if icon.is_none() && icon_mimetype.is_none(),
1713        ));
1714        Ok(())
1715    }
1716
1717    /// Tests clearing the store.
1718    #[test]
1719    fn clear() -> anyhow::Result<()> {
1720        before_each();
1721
1722        let store = TestStore::new(
1723            MockRemoteSettingsClient::default()
1724                .with_record(SuggestionProvider::Amp.record("data-1", json!([los_pollos_amp()])))
1725                .with_record(
1726                    SuggestionProvider::Amp.record("data-2", json!([good_place_eats_amp()])),
1727                )
1728                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
1729                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon()))
1730                .with_record(
1731                    SuggestionProvider::Weather
1732                        .record("weather-1", json!({ "keywords": ["abcde"], })),
1733                ),
1734        );
1735        store.ingest(SuggestIngestionConstraints::all_providers());
1736        assert!(store.count_rows("suggestions") > 0);
1737        assert!(store.count_rows("keywords") > 0);
1738        assert!(store.count_rows("keywords_i18n") > 0);
1739        assert!(store.count_rows("keywords_metrics") > 0);
1740        assert!(store.count_rows("icons") > 0);
1741
1742        store.inner.clear()?;
1743        assert!(store.count_rows("suggestions") == 0);
1744        assert!(store.count_rows("keywords") == 0);
1745        assert!(store.count_rows("keywords_i18n") == 0);
1746        assert!(store.count_rows("keywords_metrics") == 0);
1747        assert!(store.count_rows("icons") == 0);
1748
1749        Ok(())
1750    }
1751
1752    /// Tests querying suggestions.
1753    #[test]
1754    fn query() -> anyhow::Result<()> {
1755        before_each();
1756
1757        let store = TestStore::new(
1758            MockRemoteSettingsClient::default()
1759                .with_record(
1760                    SuggestionProvider::Amp.record("data-1", json!([good_place_eats_amp(),])),
1761                )
1762                .with_record(SuggestionProvider::Wikipedia.record(
1763                    "wikipedia-1",
1764                    json!([california_wiki(), caltech_wiki(), multimatch_wiki(),]),
1765                ))
1766                .with_record(
1767                    SuggestionProvider::Amo
1768                        .record("data-2", json!([relay_amo(), multimatch_amo(),])),
1769                )
1770                .with_record(SuggestionProvider::Yelp.record("data-4", json!([ramen_yelp(),])))
1771                .with_record(SuggestionProvider::Mdn.record("data-5", json!([array_mdn(),])))
1772                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon()))
1773                .with_record(SuggestionProvider::Wikipedia.icon(california_icon()))
1774                .with_record(SuggestionProvider::Wikipedia.icon(caltech_icon()))
1775                .with_record(SuggestionProvider::Yelp.icon(yelp_favicon()))
1776                .with_record(SuggestionProvider::Wikipedia.icon(multimatch_wiki_icon())),
1777        );
1778
1779        store.ingest(SuggestIngestionConstraints::all_providers());
1780
1781        assert_eq!(
1782            store.fetch_suggestions(SuggestionQuery::all_providers("")),
1783            vec![]
1784        );
1785        assert_eq!(
1786            store.fetch_suggestions(SuggestionQuery::all_providers("la")),
1787            vec![good_place_eats_suggestion("lasagna", None),]
1788        );
1789        assert_eq!(
1790            store.fetch_suggestions(SuggestionQuery::all_providers("multimatch")),
1791            vec![multimatch_amo_suggestion(), multimatch_wiki_suggestion(),]
1792        );
1793        assert_eq!(
1794            store.fetch_suggestions(SuggestionQuery::all_providers("MultiMatch")),
1795            vec![multimatch_amo_suggestion(), multimatch_wiki_suggestion(),]
1796        );
1797        assert_eq!(
1798            store.fetch_suggestions(SuggestionQuery::all_providers("multimatch").limit(1)),
1799            vec![multimatch_amo_suggestion(),],
1800        );
1801        assert_eq!(
1802            store.fetch_suggestions(SuggestionQuery::amp("la")),
1803            vec![good_place_eats_suggestion("lasagna", None)],
1804        );
1805        assert_eq!(
1806            store.fetch_suggestions(SuggestionQuery::all_providers_except(
1807                "la",
1808                SuggestionProvider::Amp
1809            )),
1810            vec![],
1811        );
1812        assert_eq!(
1813            store.fetch_suggestions(SuggestionQuery::with_providers("la", vec![])),
1814            vec![],
1815        );
1816        assert_eq!(
1817            store.fetch_suggestions(SuggestionQuery::with_providers(
1818                "cal",
1819                vec![SuggestionProvider::Amp, SuggestionProvider::Amo,]
1820            )),
1821            vec![],
1822        );
1823        assert_eq!(
1824            store.fetch_suggestions(SuggestionQuery::wikipedia("cal")),
1825            vec![
1826                california_suggestion("california"),
1827                caltech_suggestion("california"),
1828            ],
1829        );
1830        assert_eq!(
1831            store.fetch_suggestions(SuggestionQuery::wikipedia("cal").limit(1)),
1832            vec![california_suggestion("california"),],
1833        );
1834        assert_eq!(
1835            store.fetch_suggestions(SuggestionQuery::with_providers("cal", vec![])),
1836            vec![],
1837        );
1838        assert_eq!(
1839            store.fetch_suggestions(SuggestionQuery::amo("spam")),
1840            vec![relay_suggestion()],
1841        );
1842        assert_eq!(
1843            store.fetch_suggestions(SuggestionQuery::amo("masking")),
1844            vec![relay_suggestion()],
1845        );
1846        assert_eq!(
1847            store.fetch_suggestions(SuggestionQuery::amo("masking e")),
1848            vec![relay_suggestion()],
1849        );
1850        assert_eq!(
1851            store.fetch_suggestions(SuggestionQuery::amo("masking s")),
1852            vec![],
1853        );
1854        assert_eq!(
1855            store.fetch_suggestions(SuggestionQuery::with_providers(
1856                "soft",
1857                vec![SuggestionProvider::Amp, SuggestionProvider::Wikipedia]
1858            )),
1859            vec![],
1860        );
1861        assert_eq!(
1862            store.fetch_suggestions(SuggestionQuery::yelp("best spicy ramen delivery in tokyo")),
1863            vec![ramen_suggestion(
1864                "best spicy ramen delivery in tokyo",
1865                "https://www.yelp.com/search?find_desc=best+spicy+ramen+delivery&find_loc=tokyo"
1866            ),],
1867        );
1868        assert_eq!(
1869            store.fetch_suggestions(SuggestionQuery::yelp("BeSt SpIcY rAmEn DeLiVeRy In ToKyO")),
1870            vec![ramen_suggestion(
1871                "BeSt SpIcY rAmEn DeLiVeRy In ToKyO",
1872                "https://www.yelp.com/search?find_desc=BeSt+SpIcY+rAmEn+DeLiVeRy&find_loc=ToKyO"
1873            ),],
1874        );
1875        assert_eq!(
1876            store.fetch_suggestions(SuggestionQuery::yelp("best ramen delivery in tokyo")),
1877            vec![ramen_suggestion(
1878                "best ramen delivery in tokyo",
1879                "https://www.yelp.com/search?find_desc=best+ramen+delivery&find_loc=tokyo"
1880            ),],
1881        );
1882        assert_eq!(
1883            store.fetch_suggestions(SuggestionQuery::yelp(
1884                "best invalid_ramen delivery in tokyo"
1885            )),
1886            vec![],
1887        );
1888        assert_eq!(
1889            store.fetch_suggestions(SuggestionQuery::yelp("best in tokyo")),
1890            vec![],
1891        );
1892        assert_eq!(
1893            store.fetch_suggestions(SuggestionQuery::yelp("super best ramen in tokyo")),
1894            vec![ramen_suggestion(
1895                "super best ramen in tokyo",
1896                "https://www.yelp.com/search?find_desc=super+best+ramen&find_loc=tokyo"
1897            ),],
1898        );
1899        assert_eq!(
1900            store.fetch_suggestions(SuggestionQuery::yelp("invalid_best ramen in tokyo")),
1901            vec![],
1902        );
1903        assert_eq!(
1904            store.fetch_suggestions(SuggestionQuery::yelp("ramen delivery in tokyo")),
1905            vec![ramen_suggestion(
1906                "ramen delivery in tokyo",
1907                "https://www.yelp.com/search?find_desc=ramen+delivery&find_loc=tokyo"
1908            ),],
1909        );
1910        assert_eq!(
1911            store.fetch_suggestions(SuggestionQuery::yelp("ramen super delivery in tokyo")),
1912            vec![ramen_suggestion(
1913                "ramen super delivery in tokyo",
1914                "https://www.yelp.com/search?find_desc=ramen+super+delivery&find_loc=tokyo"
1915            ),],
1916        );
1917        assert_eq!(
1918            store.fetch_suggestions(SuggestionQuery::yelp("ramen invalid_delivery")),
1919            vec![ramen_suggestion(
1920                "ramen invalid_delivery",
1921                "https://www.yelp.com/search?find_desc=ramen&find_loc=invalid_delivery"
1922            )
1923            .has_location_sign(false),],
1924        );
1925        assert_eq!(
1926            store.fetch_suggestions(SuggestionQuery::yelp("ramen invalid_delivery in tokyo")),
1927            vec![ramen_suggestion(
1928                "ramen invalid_delivery in tokyo",
1929                "https://www.yelp.com/search?find_desc=ramen&find_loc=invalid_delivery+in+tokyo"
1930            )
1931            .has_location_sign(false),],
1932        );
1933        assert_eq!(
1934            store.fetch_suggestions(SuggestionQuery::yelp("ramen in tokyo")),
1935            vec![ramen_suggestion(
1936                "ramen in tokyo",
1937                "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1938            ),],
1939        );
1940        assert_eq!(
1941            store.fetch_suggestions(SuggestionQuery::yelp("ramen near tokyo")),
1942            vec![ramen_suggestion(
1943                "ramen near tokyo",
1944                "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1945            ),],
1946        );
1947        assert_eq!(
1948            store.fetch_suggestions(SuggestionQuery::yelp("ramen invalid_in tokyo")),
1949            vec![ramen_suggestion(
1950                "ramen invalid_in tokyo",
1951                "https://www.yelp.com/search?find_desc=ramen&find_loc=invalid_in+tokyo"
1952            )
1953            .has_location_sign(false),],
1954        );
1955        assert_eq!(
1956            store.fetch_suggestions(SuggestionQuery::yelp("ramen in San Francisco")),
1957            vec![ramen_suggestion(
1958                "ramen in San Francisco",
1959                "https://www.yelp.com/search?find_desc=ramen&find_loc=San+Francisco"
1960            ),],
1961        );
1962        assert_eq!(
1963            store.fetch_suggestions(SuggestionQuery::yelp("ramen in")),
1964            vec![ramen_suggestion(
1965                "ramen in",
1966                "https://www.yelp.com/search?find_desc=ramen"
1967            ),],
1968        );
1969        assert_eq!(
1970            store.fetch_suggestions(SuggestionQuery::yelp("ramen near by")),
1971            vec![ramen_suggestion(
1972                "ramen near by",
1973                "https://www.yelp.com/search?find_desc=ramen"
1974            )],
1975        );
1976        assert_eq!(
1977            store.fetch_suggestions(SuggestionQuery::yelp("ramen near me")),
1978            vec![ramen_suggestion(
1979                "ramen near me",
1980                "https://www.yelp.com/search?find_desc=ramen"
1981            )],
1982        );
1983        assert_eq!(
1984            store.fetch_suggestions(SuggestionQuery::yelp("ramen near by tokyo")),
1985            vec![ramen_suggestion(
1986                "ramen near by tokyo",
1987                "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1988            )],
1989        );
1990        assert_eq!(
1991            store.fetch_suggestions(SuggestionQuery::yelp("ramen")),
1992            vec![
1993                ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
1994                    .has_location_sign(false),
1995            ],
1996        );
1997        // Test an extremely long yelp query
1998        assert_eq!(
1999            store.fetch_suggestions(SuggestionQuery::yelp(
2000                "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789"
2001            )),
2002            vec![
2003                ramen_suggestion(
2004                    "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789",
2005                    "https://www.yelp.com/search?find_desc=012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789"
2006                ).has_location_sign(false),
2007            ],
2008        );
2009        // This query is over the limit and no suggestions should be returned
2010        assert_eq!(
2011            store.fetch_suggestions(SuggestionQuery::yelp(
2012                "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789Z"
2013            )),
2014            vec![],
2015        );
2016        assert_eq!(
2017            store.fetch_suggestions(SuggestionQuery::yelp("best delivery")),
2018            vec![],
2019        );
2020        assert_eq!(
2021            store.fetch_suggestions(SuggestionQuery::yelp("same_modifier same_modifier")),
2022            vec![],
2023        );
2024        assert_eq!(
2025            store.fetch_suggestions(SuggestionQuery::yelp("same_modifier ")),
2026            vec![],
2027        );
2028        assert_eq!(
2029            store.fetch_suggestions(SuggestionQuery::yelp("yelp ramen")),
2030            vec![
2031                ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
2032                    .has_location_sign(false),
2033            ],
2034        );
2035        assert_eq!(
2036            store.fetch_suggestions(SuggestionQuery::yelp("yelp keyword ramen")),
2037            vec![
2038                ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
2039                    .has_location_sign(false),
2040            ],
2041        );
2042        assert_eq!(
2043            store.fetch_suggestions(SuggestionQuery::yelp("ramen in tokyo yelp")),
2044            vec![ramen_suggestion(
2045                "ramen in tokyo",
2046                "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
2047            )],
2048        );
2049        assert_eq!(
2050            store.fetch_suggestions(SuggestionQuery::yelp("ramen in tokyo yelp keyword")),
2051            vec![ramen_suggestion(
2052                "ramen in tokyo",
2053                "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
2054            )],
2055        );
2056        assert_eq!(
2057            store.fetch_suggestions(SuggestionQuery::yelp("yelp ramen yelp")),
2058            vec![
2059                ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
2060                    .has_location_sign(false)
2061            ],
2062        );
2063        assert_eq!(
2064            store.fetch_suggestions(SuggestionQuery::yelp("best yelp ramen")),
2065            vec![],
2066        );
2067        assert_eq!(
2068            store.fetch_suggestions(SuggestionQuery::yelp("Spicy R")),
2069            vec![ramen_suggestion(
2070                "Spicy Ramen",
2071                "https://www.yelp.com/search?find_desc=Spicy+Ramen"
2072            )
2073            .has_location_sign(false)
2074            .subject_exact_match(false)],
2075        );
2076        assert_eq!(
2077            store.fetch_suggestions(SuggestionQuery::yelp("spi")),
2078            vec![ramen_suggestion(
2079                "spicy ramen",
2080                "https://www.yelp.com/search?find_desc=spicy+ramen"
2081            )
2082            .has_location_sign(false)
2083            .subject_exact_match(false)],
2084        );
2085        assert_eq!(
2086            store.fetch_suggestions(SuggestionQuery::yelp("BeSt             Ramen")),
2087            vec![ramen_suggestion(
2088                "BeSt Ramen",
2089                "https://www.yelp.com/search?find_desc=BeSt+Ramen"
2090            )
2091            .has_location_sign(false)],
2092        );
2093        assert_eq!(
2094            store.fetch_suggestions(SuggestionQuery::yelp("BeSt             Spicy R")),
2095            vec![ramen_suggestion(
2096                "BeSt Spicy Ramen",
2097                "https://www.yelp.com/search?find_desc=BeSt+Spicy+Ramen"
2098            )
2099            .has_location_sign(false)
2100            .subject_exact_match(false)],
2101        );
2102        assert_eq!(
2103            store.fetch_suggestions(SuggestionQuery::yelp("BeSt             R")),
2104            vec![],
2105        );
2106        assert_eq!(store.fetch_suggestions(SuggestionQuery::yelp("r")), vec![],);
2107        assert_eq!(
2108            store.fetch_suggestions(SuggestionQuery::yelp("ra")),
2109            vec![
2110                ramen_suggestion("rats", "https://www.yelp.com/search?find_desc=rats")
2111                    .has_location_sign(false)
2112                    .subject_exact_match(false)
2113            ],
2114        );
2115        assert_eq!(
2116            store.fetch_suggestions(SuggestionQuery::yelp("ram")),
2117            vec![
2118                ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
2119                    .has_location_sign(false)
2120                    .subject_exact_match(false)
2121            ],
2122        );
2123        assert_eq!(
2124            store.fetch_suggestions(SuggestionQuery::yelp("rac")),
2125            vec![
2126                ramen_suggestion("raccoon", "https://www.yelp.com/search?find_desc=raccoon")
2127                    .has_location_sign(false)
2128                    .subject_exact_match(false)
2129            ],
2130        );
2131        assert_eq!(
2132            store.fetch_suggestions(SuggestionQuery::yelp("best r")),
2133            vec![],
2134        );
2135        assert_eq!(
2136            store.fetch_suggestions(SuggestionQuery::yelp("best ra")),
2137            vec![ramen_suggestion(
2138                "best rats",
2139                "https://www.yelp.com/search?find_desc=best+rats"
2140            )
2141            .has_location_sign(false)
2142            .subject_exact_match(false)],
2143        );
2144        assert_eq!(
2145            store.fetch_suggestions(SuggestionQuery::yelp("best sp")),
2146            vec![ramen_suggestion(
2147                "best spicy ramen",
2148                "https://www.yelp.com/search?find_desc=best+spicy+ramen"
2149            )
2150            .has_location_sign(false)
2151            .subject_exact_match(false)],
2152        );
2153        assert_eq!(
2154            store.fetch_suggestions(SuggestionQuery::yelp("ramenabc")),
2155            vec![],
2156        );
2157        assert_eq!(
2158            store.fetch_suggestions(SuggestionQuery::yelp("ramenabc xyz")),
2159            vec![],
2160        );
2161        assert_eq!(
2162            store.fetch_suggestions(SuggestionQuery::yelp("best ramenabc")),
2163            vec![],
2164        );
2165        assert_eq!(
2166            store.fetch_suggestions(SuggestionQuery::yelp("bestabc ra")),
2167            vec![],
2168        );
2169        assert_eq!(
2170            store.fetch_suggestions(SuggestionQuery::yelp("bestabc ramen")),
2171            vec![],
2172        );
2173        assert_eq!(
2174            store.fetch_suggestions(SuggestionQuery::yelp("bestabc ramen xyz")),
2175            vec![],
2176        );
2177        assert_eq!(
2178            store.fetch_suggestions(SuggestionQuery::yelp("best spi ram")),
2179            vec![],
2180        );
2181        assert_eq!(
2182            store.fetch_suggestions(SuggestionQuery::yelp("bes ram")),
2183            vec![],
2184        );
2185        assert_eq!(
2186            store.fetch_suggestions(SuggestionQuery::yelp("bes ramen")),
2187            vec![],
2188        );
2189        // Test for prefix match.
2190        assert_eq!(
2191            store.fetch_suggestions(SuggestionQuery::yelp("ramen D")),
2192            vec![ramen_suggestion(
2193                "ramen Delivery",
2194                "https://www.yelp.com/search?find_desc=ramen+Delivery"
2195            )
2196            .has_location_sign(false)],
2197        );
2198        assert_eq!(
2199            store.fetch_suggestions(SuggestionQuery::yelp("ramen I")),
2200            vec![ramen_suggestion(
2201                "ramen In",
2202                "https://www.yelp.com/search?find_desc=ramen"
2203            )],
2204        );
2205        assert_eq!(
2206            store.fetch_suggestions(SuggestionQuery::yelp("ramen Y")),
2207            vec![
2208                ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
2209                    .has_location_sign(false)
2210            ],
2211        );
2212        // Prefix match is available only for last words.
2213        assert_eq!(
2214            store.fetch_suggestions(SuggestionQuery::yelp("ramen D Yelp")),
2215            vec![ramen_suggestion(
2216                "ramen D",
2217                "https://www.yelp.com/search?find_desc=ramen&find_loc=D"
2218            )
2219            .has_location_sign(false)],
2220        );
2221        assert_eq!(
2222            store.fetch_suggestions(SuggestionQuery::yelp("ramen I Tokyo")),
2223            vec![ramen_suggestion(
2224                "ramen I Tokyo",
2225                "https://www.yelp.com/search?find_desc=ramen&find_loc=I+Tokyo"
2226            )
2227            .has_location_sign(false)],
2228        );
2229        // Business subject.
2230        assert_eq!(
2231            store.fetch_suggestions(SuggestionQuery::yelp("the shop tokyo")),
2232            vec![ramen_suggestion(
2233                "the shop tokyo",
2234                "https://www.yelp.com/search?find_desc=the+shop&find_loc=tokyo"
2235            )
2236            .has_location_sign(false)
2237            .subject_type(YelpSubjectType::Business)]
2238        );
2239        assert_eq!(
2240            store.fetch_suggestions(SuggestionQuery::yelp("the sho")),
2241            vec![
2242                ramen_suggestion("the shop", "https://www.yelp.com/search?find_desc=the+shop")
2243                    .has_location_sign(false)
2244                    .subject_exact_match(false)
2245                    .subject_type(YelpSubjectType::Business)
2246            ]
2247        );
2248
2249        Ok(())
2250    }
2251
2252    // Tests querying AMP / Wikipedia
2253    #[test]
2254    fn query_with_multiple_providers_and_diff_scores() -> anyhow::Result<()> {
2255        before_each();
2256
2257        let store = TestStore::new(
2258            // Create a data set where one keyword matches multiple suggestions from each provider
2259            // where the scores are manually set.  We will test that the fetched suggestions are in
2260            // the correct order.
2261            MockRemoteSettingsClient::default()
2262                .with_record(SuggestionProvider::Amp.record(
2263                    "data-1",
2264                    json!([
2265                        los_pollos_amp().merge(json!({
2266                            "keywords": ["amp wiki match"],
2267                            "full_keywords": [("amp wiki match", 1)],
2268                            "score": 0.3,
2269                        })),
2270                        good_place_eats_amp().merge(json!({
2271                            "keywords": ["amp wiki match"],
2272                            "full_keywords": [("amp wiki match", 1)],
2273                            "score": 0.1,
2274                        })),
2275                    ]),
2276                ))
2277                .with_record(SuggestionProvider::Wikipedia.record(
2278                    "wikipedia-1",
2279                    json!([california_wiki().merge(json!({
2280                        "keywords": ["amp wiki match", "wiki match"],
2281                    })),]),
2282                ))
2283                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
2284                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon()))
2285                .with_record(SuggestionProvider::Wikipedia.icon(california_icon())),
2286        );
2287
2288        store.ingest(SuggestIngestionConstraints::all_providers());
2289        assert_eq!(
2290            store.fetch_suggestions(SuggestionQuery::all_providers("amp wiki match")),
2291            vec![
2292                los_pollos_suggestion("amp wiki match", None).with_score(0.3),
2293                // Wikipedia entries default to a 0.2 score
2294                california_suggestion("amp wiki match"),
2295                good_place_eats_suggestion("amp wiki match", None).with_score(0.1),
2296            ]
2297        );
2298        assert_eq!(
2299            store.fetch_suggestions(SuggestionQuery::all_providers("amp wiki match").limit(2)),
2300            vec![
2301                los_pollos_suggestion("amp wiki match", None).with_score(0.3),
2302                california_suggestion("amp wiki match"),
2303            ]
2304        );
2305        assert_eq!(
2306            store.fetch_suggestions(SuggestionQuery::all_providers("wiki match")),
2307            vec![california_suggestion("wiki match"),]
2308        );
2309
2310        Ok(())
2311    }
2312
2313    /// Tests ingesting malformed Remote Settings records that we understand,
2314    /// but that are missing fields, or aren't in the format we expect.
2315    #[test]
2316    fn ingest_malformed() -> anyhow::Result<()> {
2317        before_each();
2318
2319        let store = TestStore::new(
2320            MockRemoteSettingsClient::default()
2321                // Amp record without an attachment.
2322                .with_record(SuggestionProvider::Amp.empty_record("data-1"))
2323                // Wikipedia record without an attachment.
2324                .with_record(SuggestionProvider::Wikipedia.empty_record("wikipedia-1"))
2325                // Icon record without an attachment.
2326                .with_record(MockRecord {
2327                    collection: Collection::Amp,
2328                    record_type: SuggestRecordType::Icon,
2329                    id: "icon-1".to_string(),
2330                    inline_data: None,
2331                    attachment: None,
2332                })
2333                // Icon record with an ID that's not `icon-{id}`, so suggestions in
2334                // the data attachment won't be able to reference it.
2335                .with_record(MockRecord {
2336                    collection: Collection::Amp,
2337                    record_type: SuggestRecordType::Icon,
2338                    id: "bad-icon-id".to_string(),
2339                    inline_data: None,
2340                    attachment: Some(MockAttachment::Icon(MockIcon {
2341                        id: "bad-icon-id",
2342                        data: "",
2343                        mimetype: "image/png",
2344                    })),
2345                }),
2346        );
2347
2348        store.ingest(SuggestIngestionConstraints::all_providers());
2349
2350        store.read(|dao| {
2351            assert_eq!(
2352                dao.conn
2353                    .conn_ext_query_one::<i64>("SELECT count(*) FROM suggestions")?,
2354                0
2355            );
2356            assert_eq!(
2357                dao.conn
2358                    .conn_ext_query_one::<i64>("SELECT count(*) FROM icons")?,
2359                0
2360            );
2361
2362            Ok(())
2363        })?;
2364
2365        Ok(())
2366    }
2367
2368    /// Tests that we only ingest providers that we're concerned with.
2369    #[test]
2370    fn ingest_constraints_provider() -> anyhow::Result<()> {
2371        before_each();
2372
2373        let store = TestStore::new(
2374            MockRemoteSettingsClient::default()
2375                .with_record(SuggestionProvider::Amp.record("data-1", json!([los_pollos_amp()])))
2376                .with_record(SuggestionProvider::Yelp.record("yelp-1", json!([ramen_yelp()])))
2377                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon())),
2378        );
2379
2380        let constraints = SuggestIngestionConstraints {
2381            providers: Some(vec![SuggestionProvider::Amp]),
2382            ..SuggestIngestionConstraints::all_providers()
2383        };
2384        store.ingest(constraints);
2385
2386        // This should have been ingested
2387        assert_eq!(
2388            store.fetch_suggestions(SuggestionQuery::amp("lo")),
2389            vec![los_pollos_suggestion("los pollos", None)]
2390        );
2391        // This should not have been ingested, since it wasn't in the providers list
2392        assert_eq!(
2393            store.fetch_suggestions(SuggestionQuery::yelp("best ramen")),
2394            vec![]
2395        );
2396
2397        Ok(())
2398    }
2399
2400    /// Tests that records with invalid attachments are ignored
2401    #[test]
2402    fn skip_over_invalid_records() -> anyhow::Result<()> {
2403        before_each();
2404
2405        let store = TestStore::new(
2406            MockRemoteSettingsClient::default()
2407                // valid record
2408                .with_record(
2409                    SuggestionProvider::Amp.record("data-1", json!([good_place_eats_amp()])),
2410                )
2411                // This attachment is missing the `title` field and is invalid
2412                .with_record(SuggestionProvider::Amp.record(
2413                    "data-2",
2414                    json!([{
2415                            "id": 1,
2416                            "advertiser": "Los Pollos Hermanos",
2417                            "iab_category": "8 - Food & Drink",
2418                            "keywords": ["lo", "los", "los pollos"],
2419                            "url": "https://www.lph-nm.biz",
2420                            "icon": "5678",
2421                            "impression_url": "https://example.com/impression_url",
2422                            "click_url": "https://example.com/click_url",
2423                            "score": 0.3
2424                    }]),
2425                ))
2426                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon())),
2427        );
2428
2429        store.ingest(SuggestIngestionConstraints::all_providers());
2430
2431        // Test that the valid record was read
2432        assert_eq!(
2433            store.fetch_suggestions(SuggestionQuery::amp("la")),
2434            vec![good_place_eats_suggestion("lasagna", None)]
2435        );
2436        // Test that the invalid record was skipped
2437        assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("lo")), vec![]);
2438
2439        Ok(())
2440    }
2441
2442    #[test]
2443    fn query_mdn() -> anyhow::Result<()> {
2444        before_each();
2445
2446        let store = TestStore::new(
2447            MockRemoteSettingsClient::default()
2448                .with_record(SuggestionProvider::Mdn.record("mdn-1", json!([array_mdn()]))),
2449        );
2450        store.ingest(SuggestIngestionConstraints::all_providers());
2451        // prefix
2452        assert_eq!(
2453            store.fetch_suggestions(SuggestionQuery::mdn("array")),
2454            vec![array_suggestion(),]
2455        );
2456        // prefix + partial suffix
2457        assert_eq!(
2458            store.fetch_suggestions(SuggestionQuery::mdn("array java")),
2459            vec![array_suggestion(),]
2460        );
2461        // prefix + entire suffix
2462        assert_eq!(
2463            store.fetch_suggestions(SuggestionQuery::mdn("javascript array")),
2464            vec![array_suggestion(),]
2465        );
2466        // partial prefix word
2467        assert_eq!(
2468            store.fetch_suggestions(SuggestionQuery::mdn("wild")),
2469            vec![]
2470        );
2471        // single word
2472        assert_eq!(
2473            store.fetch_suggestions(SuggestionQuery::mdn("wildcard")),
2474            vec![array_suggestion()]
2475        );
2476        Ok(())
2477    }
2478
2479    #[test]
2480    fn query_no_yelp_icon_data() -> anyhow::Result<()> {
2481        before_each();
2482
2483        let store = TestStore::new(MockRemoteSettingsClient::default().with_record(
2484            SuggestionProvider::Yelp.record("yelp-1", json!([ramen_yelp()])), // Note: yelp_favicon() is missing
2485        ));
2486        store.ingest(SuggestIngestionConstraints::all_providers());
2487        assert!(matches!(
2488            store.fetch_suggestions(SuggestionQuery::yelp("ramen")).as_slice(),
2489            [Suggestion::Yelp { icon, icon_mimetype, .. }] if icon.is_none() && icon_mimetype.is_none()
2490        ));
2491
2492        Ok(())
2493    }
2494
2495    #[test]
2496    fn fetch_global_config() -> anyhow::Result<()> {
2497        before_each();
2498
2499        let store = TestStore::new(MockRemoteSettingsClient::default().with_record(MockRecord {
2500            collection: Collection::Other,
2501            record_type: SuggestRecordType::GlobalConfig,
2502            id: "configuration-1".to_string(),
2503            inline_data: Some(json!({
2504                "configuration": {
2505                    "show_less_frequently_cap": 3,
2506                },
2507            })),
2508            attachment: None,
2509        }));
2510
2511        store.ingest(SuggestIngestionConstraints::all_providers());
2512        assert_eq!(
2513            store.fetch_global_config(),
2514            SuggestGlobalConfig {
2515                show_less_frequently_cap: 3,
2516            }
2517        );
2518
2519        Ok(())
2520    }
2521
2522    #[test]
2523    fn fetch_global_config_default() -> anyhow::Result<()> {
2524        before_each();
2525
2526        let store = TestStore::new(MockRemoteSettingsClient::default());
2527        store.ingest(SuggestIngestionConstraints::all_providers());
2528        assert_eq!(
2529            store.fetch_global_config(),
2530            SuggestGlobalConfig {
2531                show_less_frequently_cap: 0,
2532            }
2533        );
2534
2535        Ok(())
2536    }
2537
2538    #[test]
2539    fn fetch_provider_config_none() -> anyhow::Result<()> {
2540        before_each();
2541
2542        let store = TestStore::new(MockRemoteSettingsClient::default());
2543        store.ingest(SuggestIngestionConstraints::all_providers());
2544        assert_eq!(store.fetch_provider_config(SuggestionProvider::Amp), None);
2545        assert_eq!(
2546            store.fetch_provider_config(SuggestionProvider::Weather),
2547            None
2548        );
2549
2550        Ok(())
2551    }
2552
2553    #[test]
2554    fn fetch_provider_config_other() -> anyhow::Result<()> {
2555        before_each();
2556
2557        let store = TestStore::new(MockRemoteSettingsClient::default().with_record(
2558            SuggestionProvider::Weather.record(
2559                "weather-1",
2560                json!({
2561                    "min_keyword_length": 3,
2562                    "score": 0.24,
2563                    "max_keyword_length": 1,
2564                    "max_keyword_word_count": 1,
2565                    "keywords": []
2566                }),
2567            ),
2568        ));
2569        store.ingest(SuggestIngestionConstraints::all_providers());
2570
2571        // Sanity-check that the weather config was ingested.
2572        assert_eq!(
2573            store.fetch_provider_config(SuggestionProvider::Weather),
2574            Some(SuggestProviderConfig::Weather {
2575                min_keyword_length: 3,
2576                score: 0.24,
2577            })
2578        );
2579
2580        // Getting the config for a different provider should return None.
2581        assert_eq!(store.fetch_provider_config(SuggestionProvider::Amp), None);
2582
2583        Ok(())
2584    }
2585
2586    #[test]
2587    fn remove_dismissed_suggestions() -> anyhow::Result<()> {
2588        before_each();
2589
2590        let store = TestStore::new(
2591            MockRemoteSettingsClient::default()
2592                .with_record(SuggestionProvider::Amp.record(
2593                    "data-1",
2594                    json!([good_place_eats_amp().merge(json!({"keywords": ["cats"]})),]),
2595                ))
2596                .with_record(SuggestionProvider::Wikipedia.record(
2597                    "wikipedia-1",
2598                    json!([california_wiki().merge(json!({"keywords": ["cats"]})),]),
2599                ))
2600                .with_record(SuggestionProvider::Amo.record(
2601                    "amo-1",
2602                    json!([relay_amo().merge(json!({"keywords": ["cats"]})),]),
2603                ))
2604                .with_record(SuggestionProvider::Mdn.record(
2605                    "mdn-1",
2606                    json!([array_mdn().merge(json!({"keywords": ["cats"]})),]),
2607                ))
2608                .with_record(SuggestionProvider::Amp.icon(good_place_eats_icon()))
2609                .with_record(SuggestionProvider::Wikipedia.icon(caltech_icon())),
2610        );
2611        store.ingest(SuggestIngestionConstraints::all_providers());
2612
2613        // A query for cats should return all suggestions
2614        let query = SuggestionQuery::all_providers("cats");
2615        let results = store.fetch_suggestions(query.clone());
2616        assert_eq!(results.len(), 4);
2617
2618        assert!(!store.inner.any_dismissed_suggestions()?);
2619
2620        for result in &results {
2621            let dismissal_key = result.dismissal_key().unwrap();
2622            assert!(!store.inner.is_dismissed_by_suggestion(result)?);
2623            assert!(!store.inner.is_dismissed_by_key(dismissal_key)?);
2624            store.inner.dismiss_by_suggestion(result)?;
2625            assert!(store.inner.is_dismissed_by_suggestion(result)?);
2626            assert!(store.inner.is_dismissed_by_key(dismissal_key)?);
2627            assert!(store.inner.any_dismissed_suggestions()?);
2628        }
2629
2630        // After dismissing the suggestions, the next query shouldn't return them
2631        assert_eq!(store.fetch_suggestions(query.clone()), vec![]);
2632
2633        // Clearing the dismissals should cause them to be returned again
2634        store.inner.clear_dismissed_suggestions()?;
2635        assert_eq!(store.fetch_suggestions(query.clone()).len(), 4);
2636
2637        for result in &results {
2638            let dismissal_key = result.dismissal_key().unwrap();
2639            assert!(!store.inner.is_dismissed_by_suggestion(result)?);
2640            assert!(!store.inner.is_dismissed_by_key(dismissal_key)?);
2641        }
2642        assert!(!store.inner.any_dismissed_suggestions()?);
2643
2644        Ok(())
2645    }
2646
2647    #[test]
2648    fn query_fakespot() -> anyhow::Result<()> {
2649        before_each();
2650
2651        let store = TestStore::new(
2652            MockRemoteSettingsClient::default()
2653                .with_record(SuggestionProvider::Fakespot.record(
2654                    "fakespot-1",
2655                    json!([snowglobe_fakespot(), simpsons_fakespot()]),
2656                ))
2657                .with_record(SuggestionProvider::Fakespot.icon(fakespot_amazon_icon())),
2658        );
2659        store.ingest(SuggestIngestionConstraints::all_providers());
2660        assert_eq!(
2661            store.fetch_suggestions(SuggestionQuery::fakespot("globe")),
2662            vec![snowglobe_suggestion(Some(FtsMatchInfo {
2663                prefix: false,
2664                stemming: false,
2665            }),)
2666            .with_fakespot_product_type_bonus(0.5)],
2667        );
2668        assert_eq!(
2669            store.fetch_suggestions(SuggestionQuery::fakespot("simpsons")),
2670            vec![simpsons_suggestion(Some(FtsMatchInfo {
2671                prefix: false,
2672                stemming: false,
2673            }),)],
2674        );
2675        // The snowglobe suggestion should come before the simpsons one, since `snow` is a partial
2676        // match on the product_type field.
2677        assert_eq!(
2678            store.fetch_suggestions(SuggestionQuery::fakespot("snow")),
2679            vec![
2680                snowglobe_suggestion(Some(FtsMatchInfo {
2681                    prefix: false,
2682                    stemming: false,
2683                }),)
2684                .with_fakespot_product_type_bonus(0.5),
2685                simpsons_suggestion(None),
2686            ],
2687        );
2688        // Test FTS by using a query where the keywords are separated in the source text
2689        assert_eq!(
2690            store.fetch_suggestions(SuggestionQuery::fakespot("simpsons snow")),
2691            vec![simpsons_suggestion(Some(FtsMatchInfo {
2692                prefix: false,
2693                stemming: false,
2694            }),)],
2695        );
2696        // Special characters should be stripped out
2697        assert_eq!(
2698            store.fetch_suggestions(SuggestionQuery::fakespot("simpsons + snow")),
2699            vec![simpsons_suggestion(Some(FtsMatchInfo {
2700                prefix: false,
2701                // This is incorrectly counted as stemming, since nothing matches the `+`
2702                // character.  TODO: fix this be improving the tokenizer in `FtsQuery`.
2703                stemming: true,
2704            }),)],
2705        );
2706
2707        Ok(())
2708    }
2709
2710    #[test]
2711    fn fakespot_keywords() -> anyhow::Result<()> {
2712        before_each();
2713
2714        let store = TestStore::new(
2715            MockRemoteSettingsClient::default()
2716                .with_record(SuggestionProvider::Fakespot.record(
2717                    "fakespot-1",
2718                    json!([
2719                        // Snow normally returns the snowglobe first.  Test using the keyword field
2720                        // to force the simpsons result first.
2721                        snowglobe_fakespot(),
2722                        simpsons_fakespot().merge(json!({"keywords": "snow"})),
2723                    ]),
2724                ))
2725                .with_record(SuggestionProvider::Fakespot.icon(fakespot_amazon_icon())),
2726        );
2727        store.ingest(SuggestIngestionConstraints::all_providers());
2728        assert_eq!(
2729            store.fetch_suggestions(SuggestionQuery::fakespot("snow")),
2730            vec![
2731                simpsons_suggestion(Some(FtsMatchInfo {
2732                    prefix: false,
2733                    stemming: false,
2734                }),)
2735                .with_fakespot_keyword_bonus(),
2736                snowglobe_suggestion(None).with_fakespot_product_type_bonus(0.5),
2737            ],
2738        );
2739        Ok(())
2740    }
2741
2742    #[test]
2743    fn fakespot_prefix_matching() -> anyhow::Result<()> {
2744        before_each();
2745
2746        let store = TestStore::new(
2747            MockRemoteSettingsClient::default()
2748                .with_record(SuggestionProvider::Fakespot.record(
2749                    "fakespot-1",
2750                    json!([snowglobe_fakespot(), simpsons_fakespot()]),
2751                ))
2752                .with_record(SuggestionProvider::Fakespot.icon(fakespot_amazon_icon())),
2753        );
2754        store.ingest(SuggestIngestionConstraints::all_providers());
2755        assert_eq!(
2756            store.fetch_suggestions(SuggestionQuery::fakespot("simp")),
2757            vec![simpsons_suggestion(Some(FtsMatchInfo {
2758                prefix: true,
2759                stemming: false,
2760            }),)],
2761        );
2762        assert_eq!(
2763            store.fetch_suggestions(SuggestionQuery::fakespot("simps")),
2764            vec![simpsons_suggestion(Some(FtsMatchInfo {
2765                prefix: true,
2766                stemming: false,
2767            }),)],
2768        );
2769        assert_eq!(
2770            store.fetch_suggestions(SuggestionQuery::fakespot("simpson")),
2771            vec![simpsons_suggestion(Some(FtsMatchInfo {
2772                prefix: false,
2773                stemming: false,
2774            }),)],
2775        );
2776
2777        Ok(())
2778    }
2779
2780    #[test]
2781    fn fakespot_updates_and_deletes() -> anyhow::Result<()> {
2782        before_each();
2783
2784        let mut store = TestStore::new(
2785            MockRemoteSettingsClient::default()
2786                .with_record(SuggestionProvider::Fakespot.record(
2787                    "fakespot-1",
2788                    json!([snowglobe_fakespot(), simpsons_fakespot()]),
2789                ))
2790                .with_record(SuggestionProvider::Fakespot.icon(fakespot_amazon_icon())),
2791        );
2792        store.ingest(SuggestIngestionConstraints::all_providers());
2793
2794        // Update the snapshot so that:
2795        //   - The Simpsons entry is deleted
2796        //   - Snow globes now use sea glass instead of glitter
2797        store
2798            .client_mut()
2799            .update_record(SuggestionProvider::Fakespot.record(
2800                "fakespot-1",
2801                json!([
2802                snowglobe_fakespot().merge(json!({"title": "Make Your Own Sea Glass Snow Globes"}))
2803            ]),
2804            ));
2805        store.ingest(SuggestIngestionConstraints::all_providers());
2806
2807        assert_eq!(
2808            store.fetch_suggestions(SuggestionQuery::fakespot("glitter")),
2809            vec![],
2810        );
2811        assert!(matches!(
2812            store.fetch_suggestions(SuggestionQuery::fakespot("sea glass")).as_slice(),
2813            [
2814                Suggestion::Fakespot { title, .. }
2815            ]
2816            if title == "Make Your Own Sea Glass Snow Globes"
2817        ));
2818
2819        assert_eq!(
2820            store.fetch_suggestions(SuggestionQuery::fakespot("simpsons")),
2821            vec![],
2822        );
2823
2824        Ok(())
2825    }
2826
2827    /// Test the pathological case where we ingest records with the same id, but from different
2828    /// collections
2829    #[test]
2830    fn same_record_id_different_collections() -> anyhow::Result<()> {
2831        before_each();
2832
2833        let mut store = TestStore::new(
2834            MockRemoteSettingsClient::default()
2835                // This record is in the fakespot-suggest-products collection
2836                .with_record(
2837                    SuggestionProvider::Fakespot
2838                        .record("fakespot-1", json!([snowglobe_fakespot()])),
2839                )
2840                // This record is in the Amp collection, but it has a fakespot record ID
2841                // for some reason.
2842                .with_record(SuggestionProvider::Amp.record("fakespot-1", json![los_pollos_amp()]))
2843                .with_record(SuggestionProvider::Amp.icon(los_pollos_icon()))
2844                .with_record(SuggestionProvider::Fakespot.icon(fakespot_amazon_icon())),
2845        );
2846        store.ingest(SuggestIngestionConstraints::all_providers());
2847        assert_eq!(
2848            store.fetch_suggestions(SuggestionQuery::fakespot("globe")),
2849            vec![snowglobe_suggestion(Some(FtsMatchInfo {
2850                prefix: false,
2851                stemming: false,
2852            }),)
2853            .with_fakespot_product_type_bonus(0.5)],
2854        );
2855        assert_eq!(
2856            store.fetch_suggestions(SuggestionQuery::amp("lo")),
2857            vec![los_pollos_suggestion("los pollos", None)],
2858        );
2859        // Test deleting one of the records
2860        store
2861            .client_mut()
2862            .delete_record(SuggestionProvider::Amp.empty_record("fakespot-1"))
2863            .delete_record(SuggestionProvider::Amp.icon(los_pollos_icon()));
2864        store.ingest(SuggestIngestionConstraints::all_providers());
2865        // FIXME(Bug 1912283): this setup currently deletes both suggestions, since
2866        // `drop_suggestions` only checks against record ID.
2867        //
2868        // assert_eq!(
2869        //     store.fetch_suggestions(SuggestionQuery::fakespot("globe")),
2870        //     vec![snowglobe_suggestion()],
2871        // );
2872        // assert_eq!(
2873        //     store.fetch_suggestions(SuggestionQuery::amp("lo")),
2874        //     vec![],
2875        // );
2876
2877        // However, we can test that the ingested records table has the correct entries
2878
2879        let record_keys = store
2880            .read(|dao| dao.get_ingested_records())
2881            .unwrap()
2882            .into_iter()
2883            .map(|r| format!("{}:{}", r.collection, r.id.as_str()))
2884            .collect::<Vec<_>>();
2885        assert_eq!(
2886            record_keys
2887                .iter()
2888                .map(String::as_str)
2889                .collect::<HashSet<_>>(),
2890            HashSet::from([
2891                "fakespot-suggest-products:fakespot-1",
2892                "fakespot-suggest-products:icon-fakespot-amazon",
2893            ]),
2894        );
2895        Ok(())
2896    }
2897
2898    #[test]
2899    fn dynamic_basic() -> anyhow::Result<()> {
2900        before_each();
2901
2902        let store = TestStore::new(
2903            MockRemoteSettingsClient::default()
2904                // A dynamic record whose attachment is a JSON object that only
2905                // contains keywords
2906                .with_record(SuggestionProvider::Dynamic.full_record(
2907                    "dynamic-0",
2908                    Some(json!({
2909                        "suggestion_type": "aaa",
2910                    })),
2911                    Some(MockAttachment::Json(json!({
2912                        "keywords": [
2913                            "aaa keyword",
2914                            "common keyword",
2915                            ["common prefix", [" aaa"]],
2916                            ["choco", ["bo", "late"]],
2917                            ["dup", ["licate 1", "licate 2"]],
2918                        ],
2919                    }))),
2920                ))
2921                // A dynamic record with a score whose attachment is a JSON
2922                // array with multiple suggestions with various properties
2923                .with_record(SuggestionProvider::Dynamic.full_record(
2924                    "dynamic-1",
2925                    Some(json!({
2926                        "suggestion_type": "bbb",
2927                        "score": 1.0,
2928                    })),
2929                    Some(MockAttachment::Json(json!([
2930                        {
2931                            "keywords": [
2932                                "bbb keyword 0",
2933                                "common keyword",
2934                                "common bbb keyword",
2935                                ["common prefix", [" bbb 0"]],
2936                            ],
2937                        },
2938                        {
2939                            "keywords": [
2940                                "bbb keyword 1",
2941                                "common keyword",
2942                                "common bbb keyword",
2943                                ["common prefix", [" bbb 1"]],
2944                            ],
2945                            "dismissal_key": "bbb-1-dismissal-key",
2946                        },
2947                        {
2948                            "keywords": [
2949                                "bbb keyword 2",
2950                                "common keyword",
2951                                "common bbb keyword",
2952                                ["common prefix", [" bbb 2"]],
2953                            ],
2954                            "data": json!("bbb-2-data"),
2955                            "dismissal_key": "bbb-2-dismissal-key",
2956                        },
2957                        {
2958                            "keywords": [
2959                                "bbb keyword 3",
2960                                "common keyword",
2961                                "common bbb keyword",
2962                                ["common prefix", [" bbb 3"]],
2963                            ],
2964                            "data": json!("bbb-3-data"),
2965                        },
2966                    ]))),
2967                )),
2968        );
2969        store.ingest(SuggestIngestionConstraints {
2970            providers: Some(vec![SuggestionProvider::Dynamic]),
2971            provider_constraints: Some(SuggestionProviderConstraints {
2972                dynamic_suggestion_types: Some(vec!["aaa".to_string(), "bbb".to_string()]),
2973                ..SuggestionProviderConstraints::default()
2974            }),
2975            ..SuggestIngestionConstraints::all_providers()
2976        });
2977
2978        // queries that shouldn't match anything
2979        let no_match_queries = vec!["aaa", "common", "common prefi", "choc", "chocolate extra"];
2980        for query in &no_match_queries {
2981            assert_eq!(
2982                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
2983                vec![],
2984            );
2985            assert_eq!(
2986                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["bbb"])),
2987                vec![],
2988            );
2989            assert_eq!(
2990                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa", "bbb"])),
2991                vec![],
2992            );
2993            assert_eq!(
2994                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa", "zzz"])),
2995                vec![],
2996            );
2997            assert_eq!(
2998                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
2999                vec![],
3000            );
3001        }
3002
3003        // queries that should match only the "aaa" suggestion
3004        let aaa_queries = [
3005            "aaa keyword",
3006            "common prefix a",
3007            "common prefix aa",
3008            "common prefix aaa",
3009            "choco",
3010            "chocob",
3011            "chocobo",
3012            "chocol",
3013            "chocolate",
3014            "dup",
3015            "dupl",
3016            "duplicate",
3017            "duplicate ",
3018            "duplicate 1",
3019            "duplicate 2",
3020        ];
3021        for query in aaa_queries {
3022            for suggestion_types in [
3023                ["aaa"].as_slice(),
3024                &["aaa", "bbb"],
3025                &["bbb", "aaa"],
3026                &["aaa", "zzz"],
3027                &["zzz", "aaa"],
3028            ] {
3029                assert_eq!(
3030                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3031                    vec![Suggestion::Dynamic {
3032                        suggestion_type: "aaa".into(),
3033                        data: None,
3034                        dismissal_key: None,
3035                        score: DEFAULT_SUGGESTION_SCORE,
3036                    }],
3037                );
3038            }
3039            assert_eq!(
3040                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["bbb"])),
3041                vec![],
3042            );
3043            assert_eq!(
3044                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3045                vec![],
3046            );
3047        }
3048
3049        // queries that should match only the "bbb 0" suggestion
3050        let bbb_0_queries = ["bbb keyword 0", "common prefix bbb 0"];
3051        for query in &bbb_0_queries {
3052            for suggestion_types in [
3053                ["bbb"].as_slice(),
3054                &["bbb", "aaa"],
3055                &["aaa", "bbb"],
3056                &["bbb", "zzz"],
3057                &["zzz", "bbb"],
3058            ] {
3059                assert_eq!(
3060                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3061                    vec![Suggestion::Dynamic {
3062                        suggestion_type: "bbb".into(),
3063                        data: None,
3064                        dismissal_key: None,
3065                        score: 1.0,
3066                    }],
3067                );
3068            }
3069            assert_eq!(
3070                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3071                vec![],
3072            );
3073            assert_eq!(
3074                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3075                vec![],
3076            );
3077        }
3078
3079        // queries that should match only the "bbb 1" suggestion
3080        let bbb_1_queries = ["bbb keyword 1", "common prefix bbb 1"];
3081        for query in &bbb_1_queries {
3082            for suggestion_types in [
3083                ["bbb"].as_slice(),
3084                &["bbb", "aaa"],
3085                &["aaa", "bbb"],
3086                &["bbb", "zzz"],
3087                &["zzz", "bbb"],
3088            ] {
3089                assert_eq!(
3090                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3091                    vec![Suggestion::Dynamic {
3092                        suggestion_type: "bbb".into(),
3093                        data: None,
3094                        dismissal_key: Some("bbb-1-dismissal-key".to_string()),
3095                        score: 1.0,
3096                    }],
3097                );
3098            }
3099            assert_eq!(
3100                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3101                vec![],
3102            );
3103            assert_eq!(
3104                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3105                vec![],
3106            );
3107        }
3108
3109        // queries that should match only the "bbb 2" suggestion
3110        let bbb_2_queries = ["bbb keyword 2", "common prefix bbb 2"];
3111        for query in &bbb_2_queries {
3112            for suggestion_types in [
3113                ["bbb"].as_slice(),
3114                &["bbb", "aaa"],
3115                &["aaa", "bbb"],
3116                &["bbb", "zzz"],
3117                &["zzz", "bbb"],
3118            ] {
3119                assert_eq!(
3120                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3121                    vec![Suggestion::Dynamic {
3122                        suggestion_type: "bbb".into(),
3123                        data: Some(json!("bbb-2-data")),
3124                        dismissal_key: Some("bbb-2-dismissal-key".to_string()),
3125                        score: 1.0,
3126                    }],
3127                );
3128            }
3129            assert_eq!(
3130                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3131                vec![],
3132            );
3133            assert_eq!(
3134                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3135                vec![],
3136            );
3137        }
3138
3139        // queries that should match only the "bbb 3" suggestion
3140        let bbb_3_queries = ["bbb keyword 3", "common prefix bbb 3"];
3141        for query in &bbb_3_queries {
3142            for suggestion_types in [
3143                ["bbb"].as_slice(),
3144                &["bbb", "aaa"],
3145                &["aaa", "bbb"],
3146                &["bbb", "zzz"],
3147                &["zzz", "bbb"],
3148            ] {
3149                assert_eq!(
3150                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3151                    vec![Suggestion::Dynamic {
3152                        suggestion_type: "bbb".into(),
3153                        data: Some(json!("bbb-3-data")),
3154                        dismissal_key: None,
3155                        score: 1.0,
3156                    }],
3157                );
3158            }
3159            assert_eq!(
3160                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3161                vec![],
3162            );
3163            assert_eq!(
3164                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3165                vec![],
3166            );
3167        }
3168
3169        // queries that should match only the "bbb" suggestions
3170        let bbb_queries = [
3171            "common bbb keyword",
3172            "common prefix b",
3173            "common prefix bb",
3174            "common prefix bbb",
3175            "common prefix bbb ",
3176        ];
3177        for query in &bbb_queries {
3178            for suggestion_types in [
3179                ["bbb"].as_slice(),
3180                &["bbb", "aaa"],
3181                &["aaa", "bbb"],
3182                &["bbb", "zzz"],
3183                &["zzz", "bbb"],
3184            ] {
3185                assert_eq!(
3186                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3187                    vec![
3188                        Suggestion::Dynamic {
3189                            suggestion_type: "bbb".into(),
3190                            data: None,
3191                            dismissal_key: None,
3192                            score: 1.0,
3193                        },
3194                        Suggestion::Dynamic {
3195                            suggestion_type: "bbb".into(),
3196                            data: None,
3197                            dismissal_key: Some("bbb-1-dismissal-key".to_string()),
3198                            score: 1.0,
3199                        },
3200                        Suggestion::Dynamic {
3201                            suggestion_type: "bbb".into(),
3202                            data: Some(json!("bbb-2-data")),
3203                            dismissal_key: Some("bbb-2-dismissal-key".to_string()),
3204                            score: 1.0,
3205                        },
3206                        Suggestion::Dynamic {
3207                            suggestion_type: "bbb".into(),
3208                            data: Some(json!("bbb-3-data")),
3209                            dismissal_key: None,
3210                            score: 1.0,
3211                        }
3212                    ],
3213                );
3214            }
3215            assert_eq!(
3216                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3217                vec![],
3218            );
3219            assert_eq!(
3220                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3221                vec![],
3222            );
3223        }
3224
3225        // queries that should match all suggestions
3226        let common_queries = ["common keyword", "common prefix", "common prefix "];
3227        for query in &common_queries {
3228            for suggestion_types in [
3229                ["aaa", "bbb"].as_slice(),
3230                &["bbb", "aaa"],
3231                &["zzz", "aaa", "bbb"],
3232                &["aaa", "zzz", "bbb"],
3233                &["aaa", "bbb", "zzz"],
3234            ] {
3235                assert_eq!(
3236                    store.fetch_suggestions(SuggestionQuery::dynamic(query, suggestion_types)),
3237                    vec![
3238                        Suggestion::Dynamic {
3239                            suggestion_type: "bbb".into(),
3240                            data: None,
3241                            dismissal_key: None,
3242                            score: 1.0,
3243                        },
3244                        Suggestion::Dynamic {
3245                            suggestion_type: "bbb".into(),
3246                            data: None,
3247                            dismissal_key: Some("bbb-1-dismissal-key".to_string()),
3248                            score: 1.0,
3249                        },
3250                        Suggestion::Dynamic {
3251                            suggestion_type: "bbb".into(),
3252                            data: Some(json!("bbb-2-data")),
3253                            dismissal_key: Some("bbb-2-dismissal-key".to_string()),
3254                            score: 1.0,
3255                        },
3256                        Suggestion::Dynamic {
3257                            suggestion_type: "bbb".into(),
3258                            data: Some(json!("bbb-3-data")),
3259                            dismissal_key: None,
3260                            score: 1.0,
3261                        },
3262                        Suggestion::Dynamic {
3263                            suggestion_type: "aaa".into(),
3264                            data: None,
3265                            dismissal_key: None,
3266                            score: DEFAULT_SUGGESTION_SCORE,
3267                        },
3268                    ],
3269                );
3270                assert_eq!(
3271                    store.fetch_suggestions(SuggestionQuery::dynamic(query, &["zzz"])),
3272                    vec![],
3273                );
3274            }
3275        }
3276
3277        Ok(())
3278    }
3279
3280    #[test]
3281    fn dynamic_same_type_in_different_records() -> anyhow::Result<()> {
3282        before_each();
3283
3284        // Make a store with the same dynamic suggestion type in three different
3285        // records.
3286        let mut store = TestStore::new(
3287            MockRemoteSettingsClient::default()
3288                // A record whose attachment is a JSON object
3289                .with_record(SuggestionProvider::Dynamic.full_record(
3290                    "dynamic-0",
3291                    Some(json!({
3292                        "suggestion_type": "aaa",
3293                    })),
3294                    Some(MockAttachment::Json(json!({
3295                        "keywords": [
3296                            "record 0 keyword",
3297                            "common keyword",
3298                            ["common prefix", [" 0"]],
3299                        ],
3300                        "data": json!("record-0-data"),
3301                    }))),
3302                ))
3303                // Another record whose attachment is a JSON object
3304                .with_record(SuggestionProvider::Dynamic.full_record(
3305                    "dynamic-1",
3306                    Some(json!({
3307                        "suggestion_type": "aaa",
3308                    })),
3309                    Some(MockAttachment::Json(json!({
3310                        "keywords": [
3311                            "record 1 keyword",
3312                            "common keyword",
3313                            ["common prefix", [" 1"]],
3314                        ],
3315                        "data": json!("record-1-data"),
3316                    }))),
3317                ))
3318                // A record whose attachment is a JSON array with some
3319                // suggestions
3320                .with_record(SuggestionProvider::Dynamic.full_record(
3321                    "dynamic-2",
3322                    Some(json!({
3323                        "suggestion_type": "aaa",
3324                    })),
3325                    Some(MockAttachment::Json(json!([
3326                        {
3327                            "keywords": [
3328                                "record 2 keyword",
3329                                "record 2 keyword 0",
3330                                "common keyword",
3331                                ["common prefix", [" 2-0"]],
3332                            ],
3333                            "data": json!("record-2-data-0"),
3334                        },
3335                        {
3336                            "keywords": [
3337                                "record 2 keyword",
3338                                "record 2 keyword 1",
3339                                "common keyword",
3340                                ["common prefix", [" 2-1"]],
3341                            ],
3342                            "data": json!("record-2-data-1"),
3343                        },
3344                    ]))),
3345                )),
3346        );
3347        store.ingest(SuggestIngestionConstraints {
3348            providers: Some(vec![SuggestionProvider::Dynamic]),
3349            provider_constraints: Some(SuggestionProviderConstraints {
3350                dynamic_suggestion_types: Some(vec!["aaa".to_string()]),
3351                ..SuggestionProviderConstraints::default()
3352            }),
3353            ..SuggestIngestionConstraints::all_providers()
3354        });
3355
3356        // queries that should match only the suggestion in record 0
3357        let record_0_queries = ["record 0 keyword", "common prefix 0"];
3358        for query in record_0_queries {
3359            assert_eq!(
3360                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3361                vec![Suggestion::Dynamic {
3362                    suggestion_type: "aaa".into(),
3363                    data: Some(json!("record-0-data")),
3364                    dismissal_key: None,
3365                    score: DEFAULT_SUGGESTION_SCORE,
3366                }],
3367            );
3368        }
3369
3370        // queries that should match only the suggestion in record 1
3371        let record_1_queries = ["record 1 keyword", "common prefix 1"];
3372        for query in record_1_queries {
3373            assert_eq!(
3374                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3375                vec![Suggestion::Dynamic {
3376                    suggestion_type: "aaa".into(),
3377                    data: Some(json!("record-1-data")),
3378                    dismissal_key: None,
3379                    score: DEFAULT_SUGGESTION_SCORE,
3380                }],
3381            );
3382        }
3383
3384        // queries that should match only the suggestions in record 2
3385        let record_2_queries = ["record 2 keyword", "common prefix 2", "common prefix 2-"];
3386        for query in record_2_queries {
3387            assert_eq!(
3388                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3389                vec![
3390                    Suggestion::Dynamic {
3391                        suggestion_type: "aaa".into(),
3392                        data: Some(json!("record-2-data-0")),
3393                        dismissal_key: None,
3394                        score: DEFAULT_SUGGESTION_SCORE,
3395                    },
3396                    Suggestion::Dynamic {
3397                        suggestion_type: "aaa".into(),
3398                        data: Some(json!("record-2-data-1")),
3399                        dismissal_key: None,
3400                        score: DEFAULT_SUGGESTION_SCORE,
3401                    },
3402                ],
3403            );
3404        }
3405
3406        // queries that should match only record 2 suggestion 0
3407        let record_2_0_queries = ["record 2 keyword 0", "common prefix 2-0"];
3408        for query in record_2_0_queries {
3409            assert_eq!(
3410                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3411                vec![Suggestion::Dynamic {
3412                    suggestion_type: "aaa".into(),
3413                    data: Some(json!("record-2-data-0")),
3414                    dismissal_key: None,
3415                    score: DEFAULT_SUGGESTION_SCORE,
3416                }],
3417            );
3418        }
3419
3420        // queries that should match only record 2 suggestion 1
3421        let record_2_1_queries = ["record 2 keyword 1", "common prefix 2-1"];
3422        for query in record_2_1_queries {
3423            assert_eq!(
3424                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3425                vec![Suggestion::Dynamic {
3426                    suggestion_type: "aaa".into(),
3427                    data: Some(json!("record-2-data-1")),
3428                    dismissal_key: None,
3429                    score: DEFAULT_SUGGESTION_SCORE,
3430                }],
3431            );
3432        }
3433
3434        // queries that should match all suggestions
3435        let common_queries = ["common keyword", "common prefix", "common prefix "];
3436        for query in common_queries {
3437            assert_eq!(
3438                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3439                vec![
3440                    Suggestion::Dynamic {
3441                        suggestion_type: "aaa".into(),
3442                        data: Some(json!("record-0-data")),
3443                        dismissal_key: None,
3444                        score: DEFAULT_SUGGESTION_SCORE,
3445                    },
3446                    Suggestion::Dynamic {
3447                        suggestion_type: "aaa".into(),
3448                        data: Some(json!("record-1-data")),
3449                        dismissal_key: None,
3450                        score: DEFAULT_SUGGESTION_SCORE,
3451                    },
3452                    Suggestion::Dynamic {
3453                        suggestion_type: "aaa".into(),
3454                        data: Some(json!("record-2-data-0")),
3455                        dismissal_key: None,
3456                        score: DEFAULT_SUGGESTION_SCORE,
3457                    },
3458                    Suggestion::Dynamic {
3459                        suggestion_type: "aaa".into(),
3460                        data: Some(json!("record-2-data-1")),
3461                        dismissal_key: None,
3462                        score: DEFAULT_SUGGESTION_SCORE,
3463                    },
3464                ],
3465            );
3466        }
3467
3468        // Delete record 0.
3469        store
3470            .client_mut()
3471            .delete_record(SuggestionProvider::Dynamic.empty_record("dynamic-0"));
3472        store.ingest(SuggestIngestionConstraints {
3473            providers: Some(vec![SuggestionProvider::Dynamic]),
3474            provider_constraints: Some(SuggestionProviderConstraints {
3475                dynamic_suggestion_types: Some(vec!["aaa".to_string()]),
3476                ..SuggestionProviderConstraints::default()
3477            }),
3478            ..SuggestIngestionConstraints::all_providers()
3479        });
3480
3481        // Keywords from record 0 should not match anything.
3482        for query in record_0_queries {
3483            assert_eq!(
3484                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3485                vec![],
3486            );
3487        }
3488
3489        // The suggestion in record 1 should remain fetchable.
3490        for query in record_1_queries {
3491            assert_eq!(
3492                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3493                vec![Suggestion::Dynamic {
3494                    suggestion_type: "aaa".into(),
3495                    data: Some(json!("record-1-data")),
3496                    dismissal_key: None,
3497                    score: DEFAULT_SUGGESTION_SCORE,
3498                }],
3499            );
3500        }
3501
3502        // The suggestions in record 2 should remain fetchable.
3503        for query in record_2_queries {
3504            assert_eq!(
3505                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3506                vec![
3507                    Suggestion::Dynamic {
3508                        suggestion_type: "aaa".into(),
3509                        data: Some(json!("record-2-data-0")),
3510                        dismissal_key: None,
3511                        score: DEFAULT_SUGGESTION_SCORE,
3512                    },
3513                    Suggestion::Dynamic {
3514                        suggestion_type: "aaa".into(),
3515                        data: Some(json!("record-2-data-1")),
3516                        dismissal_key: None,
3517                        score: DEFAULT_SUGGESTION_SCORE,
3518                    },
3519                ],
3520            );
3521        }
3522        for query in record_2_0_queries {
3523            assert_eq!(
3524                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3525                vec![Suggestion::Dynamic {
3526                    suggestion_type: "aaa".into(),
3527                    data: Some(json!("record-2-data-0")),
3528                    dismissal_key: None,
3529                    score: DEFAULT_SUGGESTION_SCORE,
3530                }],
3531            );
3532        }
3533        for query in record_2_1_queries {
3534            assert_eq!(
3535                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3536                vec![Suggestion::Dynamic {
3537                    suggestion_type: "aaa".into(),
3538                    data: Some(json!("record-2-data-1")),
3539                    dismissal_key: None,
3540                    score: DEFAULT_SUGGESTION_SCORE,
3541                }],
3542            );
3543        }
3544
3545        // All remaining suggestions should remain fetchable via the common
3546        // keywords.
3547        for query in common_queries {
3548            assert_eq!(
3549                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3550                vec![
3551                    Suggestion::Dynamic {
3552                        suggestion_type: "aaa".into(),
3553                        data: Some(json!("record-1-data")),
3554                        dismissal_key: None,
3555                        score: DEFAULT_SUGGESTION_SCORE,
3556                    },
3557                    Suggestion::Dynamic {
3558                        suggestion_type: "aaa".into(),
3559                        data: Some(json!("record-2-data-0")),
3560                        dismissal_key: None,
3561                        score: DEFAULT_SUGGESTION_SCORE,
3562                    },
3563                    Suggestion::Dynamic {
3564                        suggestion_type: "aaa".into(),
3565                        data: Some(json!("record-2-data-1")),
3566                        dismissal_key: None,
3567                        score: DEFAULT_SUGGESTION_SCORE,
3568                    },
3569                ],
3570            );
3571        }
3572
3573        // Delete record 2.
3574        store
3575            .client_mut()
3576            .delete_record(SuggestionProvider::Dynamic.empty_record("dynamic-2"));
3577        store.ingest(SuggestIngestionConstraints {
3578            providers: Some(vec![SuggestionProvider::Dynamic]),
3579            provider_constraints: Some(SuggestionProviderConstraints {
3580                dynamic_suggestion_types: Some(vec!["aaa".to_string()]),
3581                ..SuggestionProviderConstraints::default()
3582            }),
3583            ..SuggestIngestionConstraints::all_providers()
3584        });
3585
3586        // Keywords from record 0 still should not match anything.
3587        for query in record_0_queries {
3588            assert_eq!(
3589                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3590                vec![],
3591            );
3592        }
3593
3594        // The suggestion in record 1 should remain fetchable.
3595        for query in record_1_queries {
3596            assert_eq!(
3597                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3598                vec![Suggestion::Dynamic {
3599                    suggestion_type: "aaa".into(),
3600                    data: Some(json!("record-1-data")),
3601                    dismissal_key: None,
3602                    score: DEFAULT_SUGGESTION_SCORE,
3603                }],
3604            );
3605        }
3606
3607        // The suggestions in record 2 should not be fetchable.
3608        for query in record_2_queries
3609            .iter()
3610            .chain(record_2_0_queries.iter().chain(record_2_1_queries.iter()))
3611        {
3612            assert_eq!(
3613                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3614                vec![]
3615            );
3616        }
3617
3618        // The one remaining suggestion, from record 1, should remain fetchable
3619        // via the common keywords.
3620        for query in common_queries {
3621            assert_eq!(
3622                store.fetch_suggestions(SuggestionQuery::dynamic(query, &["aaa"])),
3623                vec![Suggestion::Dynamic {
3624                    suggestion_type: "aaa".into(),
3625                    data: Some(json!("record-1-data")),
3626                    dismissal_key: None,
3627                    score: DEFAULT_SUGGESTION_SCORE,
3628                },],
3629            );
3630        }
3631
3632        Ok(())
3633    }
3634
3635    #[test]
3636    fn dynamic_ingest_provider_constraints() -> anyhow::Result<()> {
3637        before_each();
3638
3639        // Create suggestions with types "aaa" and "bbb".
3640        let store = TestStore::new(
3641            MockRemoteSettingsClient::default()
3642                .with_record(SuggestionProvider::Dynamic.full_record(
3643                    "dynamic-0",
3644                    Some(json!({
3645                        "suggestion_type": "aaa",
3646                    })),
3647                    Some(MockAttachment::Json(json!({
3648                        "keywords": ["aaa keyword", "both keyword"],
3649                    }))),
3650                ))
3651                .with_record(SuggestionProvider::Dynamic.full_record(
3652                    "dynamic-1",
3653                    Some(json!({
3654                        "suggestion_type": "bbb",
3655                    })),
3656                    Some(MockAttachment::Json(json!({
3657                        "keywords": ["bbb keyword", "both keyword"],
3658                    }))),
3659                )),
3660        );
3661
3662        // Ingest but don't pass in any provider constraints. The records will
3663        // be ingested but their attachments won't be, so fetches shouldn't
3664        // return any suggestions.
3665        store.ingest(SuggestIngestionConstraints {
3666            providers: Some(vec![SuggestionProvider::Dynamic]),
3667            provider_constraints: None,
3668            ..SuggestIngestionConstraints::all_providers()
3669        });
3670
3671        let ingest_1_queries = [
3672            ("aaa keyword", vec!["aaa"]),
3673            ("aaa keyword", vec!["bbb"]),
3674            ("aaa keyword", vec!["aaa", "bbb"]),
3675            ("bbb keyword", vec!["aaa"]),
3676            ("bbb keyword", vec!["bbb"]),
3677            ("bbb keyword", vec!["aaa", "bbb"]),
3678            ("both keyword", vec!["aaa"]),
3679            ("both keyword", vec!["bbb"]),
3680            ("both keyword", vec!["aaa", "bbb"]),
3681        ];
3682        for (query, types) in &ingest_1_queries {
3683            assert_eq!(
3684                store.fetch_suggestions(SuggestionQuery::dynamic(query, types)),
3685                vec![],
3686            );
3687        }
3688
3689        // Ingest only the "bbb" suggestion. The "bbb" attachment should be
3690        // ingested, so "bbb" fetches should return the "bbb" suggestion.
3691        store.ingest(SuggestIngestionConstraints {
3692            providers: Some(vec![SuggestionProvider::Dynamic]),
3693            provider_constraints: Some(SuggestionProviderConstraints {
3694                dynamic_suggestion_types: Some(vec!["bbb".to_string()]),
3695                ..SuggestionProviderConstraints::default()
3696            }),
3697            ..SuggestIngestionConstraints::all_providers()
3698        });
3699
3700        let ingest_2_queries = [
3701            ("aaa keyword", vec!["aaa"], vec![]),
3702            ("aaa keyword", vec!["bbb"], vec![]),
3703            ("aaa keyword", vec!["aaa", "bbb"], vec![]),
3704            ("bbb keyword", vec!["aaa"], vec![]),
3705            ("bbb keyword", vec!["bbb"], vec!["bbb"]),
3706            ("bbb keyword", vec!["aaa", "bbb"], vec!["bbb"]),
3707            ("both keyword", vec!["aaa"], vec![]),
3708            ("both keyword", vec!["bbb"], vec!["bbb"]),
3709            ("both keyword", vec!["aaa", "bbb"], vec!["bbb"]),
3710        ];
3711        for (query, types, expected_types) in &ingest_2_queries {
3712            assert_eq!(
3713                store.fetch_suggestions(SuggestionQuery::dynamic(query, types)),
3714                expected_types
3715                    .iter()
3716                    .map(|t| Suggestion::Dynamic {
3717                        suggestion_type: t.to_string(),
3718                        data: None,
3719                        dismissal_key: None,
3720                        score: DEFAULT_SUGGESTION_SCORE,
3721                    })
3722                    .collect::<Vec<Suggestion>>(),
3723            );
3724        }
3725
3726        // Now ingest the "aaa" suggestion.
3727        store.ingest(SuggestIngestionConstraints {
3728            providers: Some(vec![SuggestionProvider::Dynamic]),
3729            provider_constraints: Some(SuggestionProviderConstraints {
3730                dynamic_suggestion_types: Some(vec!["aaa".to_string()]),
3731                ..SuggestionProviderConstraints::default()
3732            }),
3733            ..SuggestIngestionConstraints::all_providers()
3734        });
3735
3736        let ingest_3_queries = [
3737            ("aaa keyword", vec!["aaa"], vec!["aaa"]),
3738            ("aaa keyword", vec!["bbb"], vec![]),
3739            ("aaa keyword", vec!["aaa", "bbb"], vec!["aaa"]),
3740            ("bbb keyword", vec!["aaa"], vec![]),
3741            ("bbb keyword", vec!["bbb"], vec!["bbb"]),
3742            ("bbb keyword", vec!["aaa", "bbb"], vec!["bbb"]),
3743            ("both keyword", vec!["aaa"], vec!["aaa"]),
3744            ("both keyword", vec!["bbb"], vec!["bbb"]),
3745            ("both keyword", vec!["aaa", "bbb"], vec!["aaa", "bbb"]),
3746        ];
3747        for (query, types, expected_types) in &ingest_3_queries {
3748            assert_eq!(
3749                store.fetch_suggestions(SuggestionQuery::dynamic(query, types)),
3750                expected_types
3751                    .iter()
3752                    .map(|t| Suggestion::Dynamic {
3753                        suggestion_type: t.to_string(),
3754                        data: None,
3755                        dismissal_key: None,
3756                        score: DEFAULT_SUGGESTION_SCORE,
3757                    })
3758                    .collect::<Vec<Suggestion>>(),
3759            );
3760        }
3761
3762        Ok(())
3763    }
3764
3765    #[test]
3766    fn dynamic_ingest_new_record() -> anyhow::Result<()> {
3767        before_each();
3768
3769        // Create a dynamic suggestion and ingest it.
3770        let mut store = TestStore::new(MockRemoteSettingsClient::default().with_record(
3771            SuggestionProvider::Dynamic.full_record(
3772                "dynamic-0",
3773                Some(json!({
3774                    "suggestion_type": "aaa",
3775                })),
3776                Some(MockAttachment::Json(json!({
3777                    "keywords": ["old keyword"],
3778                }))),
3779            ),
3780        ));
3781        store.ingest(SuggestIngestionConstraints {
3782            providers: Some(vec![SuggestionProvider::Dynamic]),
3783            provider_constraints: Some(SuggestionProviderConstraints {
3784                dynamic_suggestion_types: Some(vec!["aaa".to_string()]),
3785                ..SuggestionProviderConstraints::default()
3786            }),
3787            ..SuggestIngestionConstraints::all_providers()
3788        });
3789
3790        // Add a new record of the same dynamic type.
3791        store
3792            .client_mut()
3793            .add_record(SuggestionProvider::Dynamic.full_record(
3794                "dynamic-1",
3795                Some(json!({
3796                    "suggestion_type": "aaa",
3797                })),
3798                Some(MockAttachment::Json(json!({
3799                    "keywords": ["new keyword"],
3800                }))),
3801            ));
3802
3803        // Ingest, but don't ingest the dynamic type. The store will download
3804        // the new record but shouldn't ingest its attachment.
3805        store.ingest(SuggestIngestionConstraints {
3806            providers: Some(vec![SuggestionProvider::Dynamic]),
3807            provider_constraints: None,
3808            ..SuggestIngestionConstraints::all_providers()
3809        });
3810        assert_eq!(
3811            store.fetch_suggestions(SuggestionQuery::dynamic("new keyword", &["aaa"])),
3812            vec![],
3813        );
3814
3815        // Ingest again with the dynamic type. The new record will be
3816        // unchanged, but the store should now ingest its attachment.
3817        store.ingest(SuggestIngestionConstraints {
3818            providers: Some(vec![SuggestionProvider::Dynamic]),
3819            provider_constraints: Some(SuggestionProviderConstraints {
3820                dynamic_suggestion_types: Some(vec!["aaa".to_string()]),
3821                ..SuggestionProviderConstraints::default()
3822            }),
3823            ..SuggestIngestionConstraints::all_providers()
3824        });
3825
3826        // The keyword in the new attachment should match the suggestion,
3827        // confirming that the new record's attachment was ingested.
3828        assert_eq!(
3829            store.fetch_suggestions(SuggestionQuery::dynamic("new keyword", &["aaa"])),
3830            vec![Suggestion::Dynamic {
3831                suggestion_type: "aaa".to_string(),
3832                data: None,
3833                dismissal_key: None,
3834                score: DEFAULT_SUGGESTION_SCORE,
3835            }]
3836        );
3837
3838        Ok(())
3839    }
3840
3841    #[test]
3842    fn dynamic_dismissal() -> anyhow::Result<()> {
3843        before_each();
3844
3845        let store = TestStore::new(
3846            MockRemoteSettingsClient::default()
3847                .with_record(SuggestionProvider::Dynamic.full_record(
3848                    "dynamic-0",
3849                    Some(json!({
3850                        "suggestion_type": "aaa",
3851                    })),
3852                    Some(MockAttachment::Json(json!([
3853                        {
3854                            "keywords": ["aaa"],
3855                            "dismissal_key": "dk0",
3856                        },
3857                        {
3858                            "keywords": ["aaa"],
3859                            "dismissal_key": "dk1",
3860                        },
3861                        {
3862                            "keywords": ["aaa"],
3863                        },
3864                    ]))),
3865                ))
3866                .with_record(SuggestionProvider::Dynamic.full_record(
3867                    "dynamic-1",
3868                    Some(json!({
3869                        "suggestion_type": "bbb",
3870                    })),
3871                    Some(MockAttachment::Json(json!([
3872                        {
3873                            "keywords": ["bbb"],
3874                            "dismissal_key": "dk0",
3875                        },
3876                    ]))),
3877                )),
3878        );
3879
3880        store.ingest(SuggestIngestionConstraints {
3881            providers: Some(vec![SuggestionProvider::Dynamic]),
3882            provider_constraints: Some(SuggestionProviderConstraints {
3883                dynamic_suggestion_types: Some(vec!["aaa".to_string(), "bbb".to_string()]),
3884                ..SuggestionProviderConstraints::default()
3885            }),
3886            ..SuggestIngestionConstraints::all_providers()
3887        });
3888
3889        // Make sure the suggestions are initially fetchable.
3890        assert!(!store.inner.any_dismissed_suggestions()?);
3891        let suggestions_0: Vec<Suggestion> =
3892            store.fetch_suggestions(SuggestionQuery::dynamic("aaa", &["aaa"]));
3893        let suggestions_1: Vec<Suggestion> =
3894            store.fetch_suggestions(SuggestionQuery::dynamic("bbb", &["bbb"]));
3895        assert_eq!(
3896            suggestions_0,
3897            vec![
3898                Suggestion::Dynamic {
3899                    suggestion_type: "aaa".to_string(),
3900                    data: None,
3901                    dismissal_key: Some("dk0".to_string()),
3902                    score: DEFAULT_SUGGESTION_SCORE,
3903                },
3904                Suggestion::Dynamic {
3905                    suggestion_type: "aaa".to_string(),
3906                    data: None,
3907                    dismissal_key: Some("dk1".to_string()),
3908                    score: DEFAULT_SUGGESTION_SCORE,
3909                },
3910                Suggestion::Dynamic {
3911                    suggestion_type: "aaa".to_string(),
3912                    data: None,
3913                    dismissal_key: None,
3914                    score: DEFAULT_SUGGESTION_SCORE,
3915                },
3916            ],
3917        );
3918
3919        // Dismiss the first suggestion.
3920        assert_eq!(suggestions_0[0].dismissal_key(), Some("dk0"));
3921        store.inner.dismiss_by_suggestion(&suggestions_0[0])?;
3922
3923        assert!(store.inner.any_dismissed_suggestions()?);
3924        assert!(store.inner.is_dismissed_by_suggestion(&suggestions_0[0])?);
3925        assert_eq!(
3926            store.fetch_suggestions(SuggestionQuery::dynamic("aaa", &["aaa"])),
3927            vec![
3928                Suggestion::Dynamic {
3929                    suggestion_type: "aaa".to_string(),
3930                    data: None,
3931                    dismissal_key: Some("dk1".to_string()),
3932                    score: DEFAULT_SUGGESTION_SCORE,
3933                },
3934                Suggestion::Dynamic {
3935                    suggestion_type: "aaa".to_string(),
3936                    data: None,
3937                    dismissal_key: None,
3938                    score: DEFAULT_SUGGESTION_SCORE,
3939                },
3940            ],
3941        );
3942
3943        // Dismiss the second suggestion.
3944        assert_eq!(suggestions_0[1].dismissal_key(), Some("dk1"));
3945        store.inner.dismiss_by_suggestion(&suggestions_0[1])?;
3946
3947        assert!(store.inner.is_dismissed_by_suggestion(&suggestions_0[1])?);
3948        assert_eq!(
3949            store.fetch_suggestions(SuggestionQuery::dynamic("aaa", &["aaa"])),
3950            vec![Suggestion::Dynamic {
3951                suggestion_type: "aaa".to_string(),
3952                data: None,
3953                dismissal_key: None,
3954                score: DEFAULT_SUGGESTION_SCORE,
3955            },],
3956        );
3957
3958        // Make sure the bbb suggestion hasn't been dismissed even though it
3959        // has the same key as the first aaa suggestion.
3960        assert_eq!(
3961            suggestions_1[0].dismissal_key(),
3962            suggestions_0[0].dismissal_key()
3963        );
3964        assert!(!store.inner.is_dismissed_by_suggestion(&suggestions_1[0])?);
3965        assert_eq!(
3966            store.fetch_suggestions(SuggestionQuery::dynamic("bbb", &["bbb"])),
3967            vec![Suggestion::Dynamic {
3968                suggestion_type: "bbb".to_string(),
3969                data: None,
3970                dismissal_key: Some("dk0".to_string()),
3971                score: DEFAULT_SUGGESTION_SCORE,
3972            },],
3973        );
3974
3975        // Clear dismissals. All suggestions should be fetchable again.
3976        store.inner.clear_dismissed_suggestions()?;
3977        assert_eq!(
3978            store.fetch_suggestions(SuggestionQuery::dynamic("aaa", &["aaa"])),
3979            vec![
3980                Suggestion::Dynamic {
3981                    suggestion_type: "aaa".to_string(),
3982                    data: None,
3983                    dismissal_key: Some("dk0".to_string()),
3984                    score: DEFAULT_SUGGESTION_SCORE,
3985                },
3986                Suggestion::Dynamic {
3987                    suggestion_type: "aaa".to_string(),
3988                    data: None,
3989                    dismissal_key: Some("dk1".to_string()),
3990                    score: DEFAULT_SUGGESTION_SCORE,
3991                },
3992                Suggestion::Dynamic {
3993                    suggestion_type: "aaa".to_string(),
3994                    data: None,
3995                    dismissal_key: None,
3996                    score: DEFAULT_SUGGESTION_SCORE,
3997                },
3998            ],
3999        );
4000
4001        Ok(())
4002    }
4003}