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