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