suggest/
store.rs

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