third_party/rust/suggest/src/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
   6 use std::{
   7     collections::BTreeSet,
   8     path::{Path, PathBuf},
   9     sync::Arc,
  10 };
  11
  12 use error_support::{breadcrumb, handle_error};
  13 use once_cell::sync::OnceCell;
  14 use parking_lot::Mutex;
  15 use remote_settings::{self, RemoteSettingsConfig, RemoteSettingsServer};
  16
  17 use serde::de::DeserializeOwned;
  18
  19 use crate::{
  20     config::{SuggestGlobalConfig, SuggestProviderConfig},
  21     db::{ConnectionType, SuggestDao, SuggestDb},
  22     error::Error,
  23     provider::SuggestionProvider,
  24     rs::{
  25         Client, Record, RecordRequest, SuggestAttachment, SuggestRecord, SuggestRecordId,
  26         SuggestRecordType, DEFAULT_RECORDS_TYPES, REMOTE_SETTINGS_COLLECTION,
  27     },
  28     Result, SuggestApiResult, Suggestion, SuggestionQuery,
  29 };
  30
  31 /// Builder for [SuggestStore]
  32 ///
  33 /// Using a builder is preferred to calling the constructor directly since it's harder to confuse
  34 /// the data_path and cache_path strings.
  35 pub struct SuggestStoreBuilder(Mutex<SuggestStoreBuilderInner>);
  36
  37 #[derive(Default)]
  38 struct SuggestStoreBuilderInner {
  39     data_path: Option<String>,
  40     remote_settings_server: Option<RemoteSettingsServer>,
  41     remote_settings_bucket_name: Option<String>,
  42 }
  43
  44 impl Default for SuggestStoreBuilder {
  45     fn default() -> Self {
  46         Self::new()
  47     }
  48 }
  49
  50 impl SuggestStoreBuilder {
  51     pub fn new() -> SuggestStoreBuilder {
  52         Self(Mutex::new(SuggestStoreBuilderInner::default()))
  53     }
  54
  55     pub fn data_path(self: Arc<Self>, path: String) -> Arc<Self> {
  56         self.0.lock().data_path = Some(path);
  57         self
  58     }
  59
  60     pub fn cache_path(self: Arc<Self>, _path: String) -> Arc<Self> {
  61         // We used to use this, but we're not using it anymore, just ignore the call
  62         self
  63     }
  64
  65     pub fn remote_settings_server(self: Arc<Self>, server: RemoteSettingsServer) -> Arc<Self> {
  66         self.0.lock().remote_settings_server = Some(server);
  67         self
  68     }
  69
  70     pub fn remote_settings_bucket_name(self: Arc<Self>, bucket_name: String) -> Arc<Self> {
  71         self.0.lock().remote_settings_bucket_name = Some(bucket_name);
  72         self
  73     }
  74
  75     #[handle_error(Error)]
  76     pub fn build(&self) -> SuggestApiResult<Arc<SuggestStore>> {
  77         let inner = self.0.lock();
  78         let data_path = inner
  79             .data_path
  80             .clone()
  81             .ok_or_else(|| Error::SuggestStoreBuilder("data_path not specified".to_owned()))?;
  82
  83         let remote_settings_config = RemoteSettingsConfig {
  84             server: inner.remote_settings_server.clone(),
  85             bucket_name: inner.remote_settings_bucket_name.clone(),
  86             server_url: None,
  87             collection_name: REMOTE_SETTINGS_COLLECTION.into(),
  88         };
  89         let settings_client = remote_settings::Client::new(remote_settings_config)?;
  90         Ok(Arc::new(SuggestStore {
  91             inner: SuggestStoreInner::new(data_path, settings_client),
  92         }))
  93     }
  94 }
  95
  96 /// What should be interrupted when [SuggestStore::interrupt] is called?
  97 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
  98 pub enum InterruptKind {
  99     /// Interrupt read operations like [SuggestStore::query]
 100     Read,
 101     /// Interrupt write operations.  This mostly means [SuggestStore::ingest], but
 102     /// [SuggestStore::dismiss_suggestion] may also be interrupted.
 103     Write,
 104     /// Interrupt both read and write operations,
 105     ReadWrite,
 106 }
 107
 108 /// The store is the entry point to the Suggest component. It incrementally
 109 /// downloads suggestions from the Remote Settings service, stores them in a
 110 /// local database, and returns them in response to user queries.
 111 ///
 112 /// Your application should create a single store, and manage it as a singleton.
 113 /// The store is thread-safe, and supports concurrent queries and ingests. We
 114 /// expect that your application will call [`SuggestStore::query()`] to show
 115 /// suggestions as the user types into the address bar, and periodically call
 116 /// [`SuggestStore::ingest()`] in the background to update the database with
 117 /// new suggestions from Remote Settings.
 118 ///
 119 /// For responsiveness, we recommend always calling `query()` on a worker
 120 /// thread. When the user types new input into the address bar, call
 121 /// [`SuggestStore::interrupt()`] on the main thread to cancel the query
 122 /// for the old input, and unblock the worker thread for the new query.
 123 ///
 124 /// The store keeps track of the state needed to support incremental ingestion,
 125 /// but doesn't schedule the ingestion work itself, or decide how many
 126 /// suggestions to ingest at once. This is for two reasons:
 127 ///
 128 /// 1. The primitives for scheduling background work vary between platforms, and
 129 ///    aren't available to the lower-level Rust layer. You might use an idle
 130 ///    timer on Desktop, `WorkManager` on Android, or `BGTaskScheduler` on iOS.
 131 /// 2. Ingestion constraints can change, depending on the platform and the needs
 132 ///    of your application. A mobile device on a metered connection might want
 133 ///    to request a small subset of the Suggest data and download the rest
 134 ///    later, while a desktop on a fast link might download the entire dataset
 135 ///    on the first launch.
 136 pub struct SuggestStore {
 137     inner: SuggestStoreInner<remote_settings::Client>,
 138 }
 139
 140 impl SuggestStore {
 141     /// Creates a Suggest store.
 142     #[handle_error(Error)]
 143     pub fn new(
 144         path: &str,
 145         settings_config: Option<RemoteSettingsConfig>,
 146     ) -> SuggestApiResult<Self> {
 147         let settings_client = || -> Result<_> {
 148             Ok(remote_settings::Client::new(
 149                 settings_config.unwrap_or_else(|| RemoteSettingsConfig {
 150                     server: None,
 151                     server_url: None,
 152                     bucket_name: None,
 153                     collection_name: REMOTE_SETTINGS_COLLECTION.into(),
 154                 }),
 155             )?)
 156         }()?;
 157         Ok(Self {
 158             inner: SuggestStoreInner::new(path.to_owned(), settings_client),
 159         })
 160     }
 161
 162     /// Queries the database for suggestions.
 163     #[handle_error(Error)]
 164     pub fn query(&self, query: SuggestionQuery) -> SuggestApiResult<Vec<Suggestion>> {
 165         self.inner.query(query)
 166     }
 167
 168     /// Dismiss a suggestion
 169     ///
 170     /// Dismissed suggestions will not be returned again
 171     ///
 172     /// In the case of AMP suggestions this should be the raw URL.
 173     #[handle_error(Error)]
 174     pub fn dismiss_suggestion(&self, suggestion_url: String) -> SuggestApiResult<()> {
 175         self.inner.dismiss_suggestion(suggestion_url)
 176     }
 177
 178     /// Clear dismissed suggestions
 179     #[handle_error(Error)]
 180     pub fn clear_dismissed_suggestions(&self) -> SuggestApiResult<()> {
 181         self.inner.clear_dismissed_suggestions()
 182     }
 183
 184     /// Interrupts any ongoing queries.
 185     ///
 186     /// This should be called when the user types new input into the address
 187     /// bar, to ensure that they see fresh suggestions as they type. This
 188     /// method does not interrupt any ongoing ingests.
 189     pub fn interrupt(&self, kind: Option<InterruptKind>) {
 190         self.inner.interrupt(kind)
 191     }
 192
 193     /// Ingests new suggestions from Remote Settings.
 194     #[handle_error(Error)]
 195     pub fn ingest(&self, constraints: SuggestIngestionConstraints) -> SuggestApiResult<()> {
 196         self.inner.ingest(constraints)
 197     }
 198
 199     /// Removes all content from the database.
 200     #[handle_error(Error)]
 201     pub fn clear(&self) -> SuggestApiResult<()> {
 202         self.inner.clear()
 203     }
 204
 205     // Returns global Suggest configuration data.
 206     #[handle_error(Error)]
 207     pub fn fetch_global_config(&self) -> SuggestApiResult<SuggestGlobalConfig> {
 208         self.inner.fetch_global_config()
 209     }
 210
 211     // Returns per-provider Suggest configuration data.
 212     #[handle_error(Error)]
 213     pub fn fetch_provider_config(
 214         &self,
 215         provider: SuggestionProvider,
 216     ) -> SuggestApiResult<Option<SuggestProviderConfig>> {
 217         self.inner.fetch_provider_config(provider)
 218     }
 219 }
 220
 221 /// Constraints limit which suggestions to ingest from Remote Settings.
 222 #[derive(Clone, Default, Debug)]
 223 pub struct SuggestIngestionConstraints {
 224     /// The approximate maximum number of suggestions to ingest. Set to [`None`]
 225     /// for "no limit".
 226     ///
 227     /// Because of how suggestions are partitioned in Remote Settings, this is a
 228     /// soft limit, and the store might ingest more than requested.
 229     pub max_suggestions: Option<u64>,
 230     pub providers: Option<Vec<SuggestionProvider>>,
 231     /// Only run ingestion if the table `suggestions` is empty
 232     pub empty_only: bool,
 233 }
 234
 235 /// The implementation of the store. This is generic over the Remote Settings
 236 /// client, and is split out from the concrete [`SuggestStore`] for testing
 237 /// with a mock client.
 238 pub(crate) struct SuggestStoreInner<S> {
 239     /// Path to the persistent SQL database.
 240     ///
 241     /// This stores things that should persist when the user clears their cache.
 242     /// It's not currently used because not all consumers pass this in yet.
 243     #[allow(unused)]
 244     data_path: PathBuf,
 245     dbs: OnceCell<SuggestStoreDbs>,
 246     settings_client: S,
 247 }
 248
 249 impl<S> SuggestStoreInner<S> {
 250     pub fn new(data_path: impl Into<PathBuf>, settings_client: S) -> Self {
 251         Self {
 252             data_path: data_path.into(),
 253             dbs: OnceCell::new(),
 254             settings_client,
 255         }
 256     }
 257
 258     /// Returns this store's database connections, initializing them if
 259     /// they're not already open.
 260     fn dbs(&self) -> Result<&SuggestStoreDbs> {
 261         self.dbs
 262             .get_or_try_init(|| SuggestStoreDbs::open(&self.data_path))
 263     }
 264
 265     fn query(&self, query: SuggestionQuery) -> Result<Vec<Suggestion>> {
 266         if query.keyword.is_empty() || query.providers.is_empty() {
 267             return Ok(Vec::new());
 268         }
 269         self.dbs()?.reader.read(|dao| dao.fetch_suggestions(&query))
 270     }
 271
 272     fn dismiss_suggestion(&self, suggestion_url: String) -> Result<()> {
 273         self.dbs()?
 274             .writer
 275             .write(|dao| dao.insert_dismissal(&suggestion_url))
 276     }
 277
 278     fn clear_dismissed_suggestions(&self) -> Result<()> {
 279         self.dbs()?.writer.write(|dao| dao.clear_dismissals())?;
 280         Ok(())
 281     }
 282
 283     fn interrupt(&self, kind: Option<InterruptKind>) {
 284         if let Some(dbs) = self.dbs.get() {
 285             // Only interrupt if the databases are already open.
 286             match kind.unwrap_or(InterruptKind::Read) {
 287                 InterruptKind::Read => {
 288                     dbs.reader.interrupt_handle.interrupt();
 289                 }
 290                 InterruptKind::Write => {
 291                     dbs.writer.interrupt_handle.interrupt();
 292                 }
 293                 InterruptKind::ReadWrite => {
 294                     dbs.reader.interrupt_handle.interrupt();
 295                     dbs.writer.interrupt_handle.interrupt();
 296                 }
 297             }
 298         }
 299     }
 300
 301     fn clear(&self) -> Result<()> {
 302         self.dbs()?.writer.write(|dao| dao.clear())
 303     }
 304
 305     pub fn fetch_global_config(&self) -> Result<SuggestGlobalConfig> {
 306         self.dbs()?.reader.read(|dao| dao.get_global_config())
 307     }
 308
 309     pub fn fetch_provider_config(
 310         &self,
 311         provider: SuggestionProvider,
 312     ) -> Result<Option<SuggestProviderConfig>> {
 313         self.dbs()?
 314             .reader
 315             .read(|dao| dao.get_provider_config(provider))
 316     }
 317 }
 318
 319 impl<S> SuggestStoreInner<S>
 320 where
 321     S: Client,
 322 {
 323     pub fn ingest(&self, constraints: SuggestIngestionConstraints) -> Result<()> {
 324         breadcrumb!("Ingestion starting");
 325         let writer = &self.dbs()?.writer;
 326         if constraints.empty_only && !writer.read(|dao| dao.suggestions_table_empty())? {
 327             return Ok(());
 328         }
 329
 330         // use std::collections::BTreeSet;
 331         let ingest_record_types = if let Some(rt) = &constraints.providers {
 332             rt.iter()
 333                 .flat_map(|x| x.records_for_provider())
 334                 .collect::<BTreeSet<_>>()
 335                 .into_iter()
 336                 .collect()
 337         } else {
 338             DEFAULT_RECORDS_TYPES.to_vec()
 339         };
 340
 341         // Handle ingestion inside single write scope
 342         let mut write_scope = writer.write_scope()?;
 343         for ingest_record_type in ingest_record_types {
 344             breadcrumb!("Ingesting {ingest_record_type}");
 345             write_scope
 346                 .write(|dao| self.ingest_records_by_type(ingest_record_type, dao, &constraints))?;
 347             write_scope.err_if_interrupted()?;
 348         }
 349         breadcrumb!("Ingestion complete");
 350
 351         Ok(())
 352     }
 353
 354     fn ingest_records_by_type(
 355         &self,
 356         ingest_record_type: SuggestRecordType,
 357         dao: &mut SuggestDao,
 358         constraints: &SuggestIngestionConstraints,
 359     ) -> Result<()> {
 360         let request = RecordRequest {
 361             record_type: Some(ingest_record_type.to_string()),
 362             last_modified: dao
 363                 .get_meta::<u64>(ingest_record_type.last_ingest_meta_key().as_str())?,
 364             limit: constraints.max_suggestions,
 365         };
 366
 367         let records = self.settings_client.get_records(request)?;
 368         self.ingest_records(&ingest_record_type.last_ingest_meta_key(), dao, &records)?;
 369         Ok(())
 370     }
 371
 372     fn ingest_records(
 373         &self,
 374         last_ingest_key: &str,
 375         dao: &mut SuggestDao,
 376         records: &[Record],
 377     ) -> Result<()> {
 378         for record in records {
 379             let record_id = SuggestRecordId::from(&record.id);
 380             if record.deleted {
 381                 // If the entire record was deleted, drop all its suggestions
 382                 // and advance the last ingest time.
 383                 dao.handle_deleted_record(last_ingest_key, record)?;
 384                 continue;
 385             }
 386             let Ok(fields) =
 387                 serde_json::from_value(serde_json::Value::Object(record.fields.clone()))
 388             else {
 389                 // We don't recognize this record's type, so we don't know how
 390                 // to ingest its suggestions. Skip processing this record.
 391                 continue;
 392             };
 393
 394             match fields {
 395                 SuggestRecord::AmpWikipedia => {
 396                     self.ingest_attachment(
 397                         // TODO: Currently re-creating the last_ingest_key because using last_ingest_meta
 398                         // breaks the tests (particularly the unparsable functionality). So, keeping
 399                         // a direct reference until we remove the "unparsable" functionality.
 400                         &SuggestRecordType::AmpWikipedia.last_ingest_meta_key(),
 401                         dao,
 402                         record,
 403                         |dao, record_id, suggestions| {
 404                             dao.insert_amp_wikipedia_suggestions(record_id, suggestions)
 405                         },
 406                     )?;
 407                 }
 408                 SuggestRecord::AmpMobile => {
 409                     self.ingest_attachment(
 410                         &SuggestRecordType::AmpMobile.last_ingest_meta_key(),
 411                         dao,
 412                         record,
 413                         |dao, record_id, suggestions| {
 414                             dao.insert_amp_mobile_suggestions(record_id, suggestions)
 415                         },
 416                     )?;
 417                 }
 418                 SuggestRecord::Icon => {
 419                     let (Some(icon_id), Some(attachment)) =
 420                         (record_id.as_icon_id(), record.attachment.as_ref())
 421                     else {
 422                         // An icon record should have an icon ID and an
 423                         // attachment. Icons that don't have these are
 424                         // malformed, so skip to the next record.
 425                         dao.put_last_ingest_if_newer(
 426                             &SuggestRecordType::Icon.last_ingest_meta_key(),
 427                             record.last_modified,
 428                         )?;
 429                         continue;
 430                     };
 431                     let data = record.require_attachment_data()?;
 432                     dao.put_icon(icon_id, data, &attachment.mimetype)?;
 433                     dao.handle_ingested_record(
 434                         &SuggestRecordType::Icon.last_ingest_meta_key(),
 435                         record,
 436                     )?;
 437                 }
 438                 SuggestRecord::Amo => {
 439                     self.ingest_attachment(
 440                         &SuggestRecordType::Amo.last_ingest_meta_key(),
 441                         dao,
 442                         record,
 443                         |dao, record_id, suggestions| {
 444                             dao.insert_amo_suggestions(record_id, suggestions)
 445                         },
 446                     )?;
 447                 }
 448                 SuggestRecord::Pocket => {
 449                     self.ingest_attachment(
 450                         &SuggestRecordType::Pocket.last_ingest_meta_key(),
 451                         dao,
 452                         record,
 453                         |dao, record_id, suggestions| {
 454                             dao.insert_pocket_suggestions(record_id, suggestions)
 455                         },
 456                     )?;
 457                 }
 458                 SuggestRecord::Yelp => {
 459                     self.ingest_attachment(
 460                         &SuggestRecordType::Yelp.last_ingest_meta_key(),
 461                         dao,
 462                         record,
 463                         |dao, record_id, suggestions| match suggestions.first() {
 464                             Some(suggestion) => dao.insert_yelp_suggestions(record_id, suggestion),
 465                             None => Ok(()),
 466                         },
 467                     )?;
 468                 }
 469                 SuggestRecord::Mdn => {
 470                     self.ingest_attachment(
 471                         &SuggestRecordType::Mdn.last_ingest_meta_key(),
 472                         dao,
 473                         record,
 474                         |dao, record_id, suggestions| {
 475                             dao.insert_mdn_suggestions(record_id, suggestions)
 476                         },
 477                     )?;
 478                 }
 479                 SuggestRecord::Weather(data) => {
 480                     self.ingest_record(
 481                         &SuggestRecordType::Weather.last_ingest_meta_key(),
 482                         dao,
 483                         record,
 484                         |dao, record_id| dao.insert_weather_data(record_id, &data),
 485                     )?;
 486                 }
 487                 SuggestRecord::GlobalConfig(config) => {
 488                     self.ingest_record(
 489                         &SuggestRecordType::GlobalConfig.last_ingest_meta_key(),
 490                         dao,
 491                         record,
 492                         |dao, _| dao.put_global_config(&SuggestGlobalConfig::from(&config)),
 493                     )?;
 494                 }
 495             }
 496         }
 497         Ok(())
 498     }
 499
 500     fn ingest_record(
 501         &self,
 502         last_ingest_key: &str,
 503         dao: &mut SuggestDao,
 504         record: &Record,
 505         ingestion_handler: impl FnOnce(&mut SuggestDao<'_>, &SuggestRecordId) -> Result<()>,
 506     ) -> Result<()> {
 507         let record_id = SuggestRecordId::from(&record.id);
 508
 509         // Drop any data that we previously ingested from this record.
 510         // Suggestions in particular don't have a stable identifier, and
 511         // determining which suggestions in the record actually changed is
 512         // more complicated than dropping and re-ingesting all of them.
 513         dao.drop_suggestions(&record_id)?;
 514
 515         // Ingest (or re-ingest) all data in the record.
 516         ingestion_handler(dao, &record_id)?;
 517
 518         dao.handle_ingested_record(last_ingest_key, record)
 519     }
 520
 521     fn ingest_attachment<T>(
 522         &self,
 523         last_ingest_key: &str,
 524         dao: &mut SuggestDao,
 525         record: &Record,
 526         ingestion_handler: impl FnOnce(&mut SuggestDao<'_>, &SuggestRecordId, &[T]) -> Result<()>,
 527     ) -> Result<()>
 528     where
 529         T: DeserializeOwned,
 530     {
 531         if record.attachment.is_none() {
 532             // This method should be called only when a record is expected to
 533             // have an attachment. If it doesn't have one, it's malformed, so
 534             // skip to the next record.
 535             dao.put_last_ingest_if_newer(last_ingest_key, record.last_modified)?;
 536             return Ok(());
 537         };
 538
 539         let attachment_data = record.require_attachment_data()?;
 540         match serde_json::from_slice::<SuggestAttachment<T>>(attachment_data) {
 541             Ok(attachment) => self.ingest_record(last_ingest_key, dao, record, |dao, record_id| {
 542                 ingestion_handler(dao, record_id, attachment.suggestions())
 543             }),
 544             // If the attachment doesn't match our expected schema, just skip it.  It's possible
 545             // that we're using an older version.  If so, we'll get the data when we re-ingest
 546             // after updating the schema.
 547             Err(_) => Ok(()),
 548         }
 549     }
 550 }
 551
 552 #[cfg(feature = "benchmark_api")]
 553 impl<S> SuggestStoreInner<S>
 554 where
 555     S: Client,
 556 {
 557     pub fn into_settings_client(self) -> S {
 558         self.settings_client
 559     }
 560
 561     pub fn ensure_db_initialized(&self) {
 562         self.dbs().unwrap();
 563     }
 564
 565     pub fn force_reingest(&self, ingest_record_type: SuggestRecordType) {
 566         // To force a re-ingestion, we're going to ingest all records then forget the last
 567         // ingestion time.
 568         self.benchmark_ingest_records_by_type(ingest_record_type);
 569         let writer = &self.dbs().unwrap().writer;
 570         writer
 571             .write(|dao| dao.clear_meta(ingest_record_type.last_ingest_meta_key().as_str()))
 572             .unwrap();
 573     }
 574
 575     pub fn benchmark_ingest_records_by_type(&self, ingest_record_type: SuggestRecordType) {
 576         let writer = &self.dbs().unwrap().writer;
 577         writer
 578             .write(|dao| {
 579                 dao.clear_meta(ingest_record_type.last_ingest_meta_key().as_str())?;
 580                 self.ingest_records_by_type(
 581                     ingest_record_type,
 582                     dao,
 583                     &SuggestIngestionConstraints::default(),
 584                 )
 585             })
 586             .unwrap()
 587     }
 588
 589     pub fn table_row_counts(&self) -> Vec<(String, u32)> {
 590         use sql_support::ConnExt;
 591
 592         // Note: since this is just used for debugging, use unwrap to simplify the error handling.
 593         let reader = &self.dbs().unwrap().reader;
 594         let conn = reader.conn.lock();
 595         let table_names: Vec<String> = conn
 596             .query_rows_and_then(
 597                 "SELECT name FROM sqlite_master where type = 'table'",
 598                 (),
 599                 |row| row.get(0),
 600             )
 601             .unwrap();
 602         let mut table_names_with_counts: Vec<(String, u32)> = table_names
 603             .into_iter()
 604             .map(|name| {
 605                 let count: u32 = conn
 606                     .query_one(&format!("SELECT COUNT(*) FROM {name}"))
 607                     .unwrap();
 608                 (name, count)
 609             })
 610             .collect();
 611         table_names_with_counts.sort_by(|a, b| (b.1.cmp(&a.1)));
 612         table_names_with_counts
 613     }
 614
 615     pub fn db_size(&self) -> usize {
 616         use sql_support::ConnExt;
 617
 618         let reader = &self.dbs().unwrap().reader;
 619         let conn = reader.conn.lock();
 620         conn.query_one("SELECT page_size * page_count FROM pragma_page_count(), pragma_page_size()")
 621             .unwrap()
 622     }
 623 }
 624
 625 /// Holds a store's open connections to the Suggest database.
 626 struct SuggestStoreDbs {
 627     /// A read-write connection used to update the database with new data.
 628     writer: SuggestDb,
 629     /// A read-only connection used to query the database.
 630     reader: SuggestDb,
 631 }
 632
 633 impl SuggestStoreDbs {
 634     fn open(path: &Path) -> Result<Self> {
 635         // Order is important here: the writer must be opened first, so that it
 636         // can set up the database and run any migrations.
 637         let writer = SuggestDb::open(path, ConnectionType::ReadWrite)?;
 638         let reader = SuggestDb::open(path, ConnectionType::ReadOnly)?;
 639         Ok(Self { writer, reader })
 640     }
 641 }
 642
 643 #[cfg(test)]
 644 mod tests {
 645     use super::*;
 646
 647     use std::sync::atomic::{AtomicUsize, Ordering};
 648
 649     use parking_lot::Once;
 650     use serde_json::json;
 651     use sql_support::ConnExt;
 652
 653     use crate::{testing::*, SuggestionProvider};
 654
 655     /// In-memory Suggest store for testing
 656     struct TestStore {
 657         pub inner: SuggestStoreInner<MockRemoteSettingsClient>,
 658     }
 659
 660     impl TestStore {
 661         fn new(client: MockRemoteSettingsClient) -> Self {
 662             static COUNTER: AtomicUsize = AtomicUsize::new(0);
 663             let db_path = format!(
 664                 "file:test_store_data_{}?mode=memory&cache=shared",
 665                 COUNTER.fetch_add(1, Ordering::Relaxed),
 666             );
 667             Self {
 668                 inner: SuggestStoreInner::new(db_path, client),
 669             }
 670         }
 671
 672         fn replace_client(&mut self, client: MockRemoteSettingsClient) {
 673             self.inner.settings_client = client;
 674         }
 675
 676         fn last_modified_timestamp(&self) -> u64 {
 677             self.inner.settings_client.last_modified_timestamp
 678         }
 679
 680         fn read<T>(&self, op: impl FnOnce(&SuggestDao) -> Result<T>) -> Result<T> {
 681             self.inner.dbs().unwrap().reader.read(op)
 682         }
 683
 684         fn write<T>(&self, op: impl FnOnce(&mut SuggestDao) -> Result<T>) -> Result<T> {
 685             self.inner.dbs().unwrap().writer.write(op)
 686         }
 687
 688         fn count_rows(&self, table_name: &str) -> u64 {
 689             let sql = format!("SELECT count(*) FROM {table_name}");
 690             self.read(|dao| Ok(dao.conn.query_one(&sql)?))
 691                 .unwrap_or_else(|e| panic!("SQL error in count: {e}"))
 692         }
 693
 694         fn ingest(&self, constraints: SuggestIngestionConstraints) {
 695             self.inner.ingest(constraints).unwrap();
 696         }
 697
 698         fn fetch_suggestions(&self, query: SuggestionQuery) -> Vec<Suggestion> {
 699             self.inner
 700                 .dbs()
 701                 .unwrap()
 702                 .reader
 703                 .read(|dao| Ok(dao.fetch_suggestions(&query).unwrap()))
 704                 .unwrap()
 705         }
 706
 707         pub fn fetch_global_config(&self) -> SuggestGlobalConfig {
 708             self.inner
 709                 .fetch_global_config()
 710                 .expect("Error fetching global config")
 711         }
 712
 713         pub fn fetch_provider_config(
 714             &self,
 715             provider: SuggestionProvider,
 716         ) -> Option<SuggestProviderConfig> {
 717             self.inner
 718                 .fetch_provider_config(provider)
 719                 .expect("Error fetching provider config")
 720         }
 721     }
 722
 723     fn before_each() {
 724         static ONCE: Once = Once::new();
 725         ONCE.call_once(|| {
 726             env_logger::init();
 727         });
 728     }
 729
 730     /// Tests that `SuggestStore` is usable with UniFFI, which requires exposed
 731     /// interfaces to be `Send` and `Sync`.
 732     #[test]
 733     fn is_thread_safe() {
 734         before_each();
 735
 736         fn is_send_sync<T: Send + Sync>() {}
 737         is_send_sync::<SuggestStore>();
 738     }
 739
 740     /// Tests ingesting suggestions into an empty database.
 741     #[test]
 742     fn ingest_suggestions() -> anyhow::Result<()> {
 743         before_each();
 744
 745         let store = TestStore::new(
 746             MockRemoteSettingsClient::default()
 747                 .with_record("data", "1234", json![los_pollos_amp()])
 748                 .with_icon(los_pollos_icon()),
 749         );
 750         store.ingest(SuggestIngestionConstraints::default());
 751         assert_eq!(
 752             store.fetch_suggestions(SuggestionQuery::amp("lo")),
 753             vec![los_pollos_suggestion("los")],
 754         );
 755         Ok(())
 756     }
 757
 758     /// Tests ingesting suggestions into an empty database.
 759     #[test]
 760     fn ingest_empty_only() -> anyhow::Result<()> {
 761         before_each();
 762
 763         let mut store = TestStore::new(MockRemoteSettingsClient::default().with_record(
 764             "data",
 765             "1234",
 766             json![los_pollos_amp()],
 767         ));
 768         // suggestions_table_empty returns true before the ingestion is complete
 769         assert!(store.read(|dao| dao.suggestions_table_empty())?);
 770         // This ingestion should run, since the DB is empty
 771         store.ingest(SuggestIngestionConstraints {
 772             empty_only: true,
 773             ..SuggestIngestionConstraints::default()
 774         });
 775         // suggestions_table_empty returns false after the ingestion is complete
 776         assert!(!store.read(|dao| dao.suggestions_table_empty())?);
 777
 778         // This ingestion should not run since the DB is no longer empty
 779         store.replace_client(MockRemoteSettingsClient::default().with_record(
 780             "data",
 781             "1234",
 782             json!([los_pollos_amp(), good_place_eats_amp()]),
 783         ));
 784         store.ingest(SuggestIngestionConstraints {
 785             empty_only: true,
 786             ..SuggestIngestionConstraints::default()
 787         });
 788         // "la" should not match the good place eats suggestion, since that should not have been
 789         // ingested.
 790         assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("la")), vec![]);
 791
 792         Ok(())
 793     }
 794
 795     /// Tests ingesting suggestions with icons.
 796     #[test]
 797     fn ingest_amp_icons() -> anyhow::Result<()> {
 798         before_each();
 799
 800         let store = TestStore::new(
 801             MockRemoteSettingsClient::default()
 802                 .with_record(
 803                     "data",
 804                     "1234",
 805                     json!([los_pollos_amp(), good_place_eats_amp()]),
 806                 )
 807                 .with_icon(los_pollos_icon())
 808                 .with_icon(good_place_eats_icon()),
 809         );
 810         // This ingestion should run, since the DB is empty
 811         store.ingest(SuggestIngestionConstraints::default());
 812
 813         assert_eq!(
 814             store.fetch_suggestions(SuggestionQuery::amp("lo")),
 815             vec![los_pollos_suggestion("los")]
 816         );
 817         assert_eq!(
 818             store.fetch_suggestions(SuggestionQuery::amp("la")),
 819             vec![good_place_eats_suggestion("lasagna")]
 820         );
 821
 822         Ok(())
 823     }
 824
 825     #[test]
 826     fn ingest_full_keywords() -> anyhow::Result<()> {
 827         before_each();
 828
 829         let store = TestStore::new(MockRemoteSettingsClient::default()
 830             .with_record("data", "1234", json!([
 831                 // AMP attachment with full keyword data
 832                 los_pollos_amp().merge(json!({
 833                     "keywords": ["lo", "los", "los p", "los pollos", "los pollos h", "los pollos hermanos"],
 834                     "full_keywords": [
 835                         // Full keyword for the first 4 keywords
 836                         ("los pollos", 4),
 837                         // Full keyword for the next 2 keywords
 838                         ("los pollos hermanos (restaurant)", 2),
 839                     ],
 840                 })),
 841                 // AMP attachment without full keyword data
 842                 good_place_eats_amp(),
 843                 // Wikipedia attachment with full keyword data.  We should ignore the full
 844                 // keyword data for Wikipedia suggestions
 845                 california_wiki(),
 846                 // california_wiki().merge(json!({
 847                 //     "keywords": ["cal", "cali", "california"],
 848                 //     "full_keywords": [("california institute of technology", 3)],
 849                 // })),
 850             ]))
 851             .with_record("amp-mobile-suggestions", "2468", json!([
 852                 // Amp mobile attachment with full keyword data
 853                 a1a_amp_mobile().merge(json!({
 854                     "keywords": ["a1a", "ca", "car", "car wash"],
 855                     "full_keywords": [
 856                         ("A1A Car Wash", 1),
 857                         ("car wash", 3),
 858                     ],
 859                 })),
 860             ]))
 861             .with_icon(los_pollos_icon())
 862             .with_icon(good_place_eats_icon())
 863             .with_icon(california_icon())
 864         );
 865         store.ingest(SuggestIngestionConstraints::default());
 866
 867         assert_eq!(
 868             store.fetch_suggestions(SuggestionQuery::amp("lo")),
 869             // This keyword comes from the provided full_keywords list
 870             vec![los_pollos_suggestion("los pollos")],
 871         );
 872
 873         assert_eq!(
 874             store.fetch_suggestions(SuggestionQuery::amp("la")),
 875             // Good place eats did not have full keywords, so this one is calculated with the
 876             // keywords.rs code
 877             vec![good_place_eats_suggestion("lasagna")],
 878         );
 879
 880         assert_eq!(
 881             store.fetch_suggestions(SuggestionQuery::wikipedia("cal")),
 882             // Even though this had a full_keywords field, we should ignore it since it's a
 883             // wikipedia suggestion and use the keywords.rs code instead
 884             vec![california_suggestion("california")],
 885         );
 886
 887         assert_eq!(
 888             store.fetch_suggestions(SuggestionQuery::amp_mobile("a1a")),
 889             // This keyword comes from the provided full_keywords list.
 890             vec![a1a_suggestion("A1A Car Wash")],
 891         );
 892
 893         Ok(())
 894     }
 895
 896     /// Tests ingesting a data attachment containing a single suggestion,
 897     /// instead of an array of suggestions.
 898     #[test]
 899     fn ingest_one_suggestion_in_data_attachment() -> anyhow::Result<()> {
 900         before_each();
 901
 902         let store = TestStore::new(
 903             MockRemoteSettingsClient::default()
 904                 // This record contains just one JSON object, rather than an array of them
 905                 .with_record("data", "1234", los_pollos_amp())
 906                 .with_icon(los_pollos_icon()),
 907         );
 908         store.ingest(SuggestIngestionConstraints::default());
 909         assert_eq!(
 910             store.fetch_suggestions(SuggestionQuery::amp("lo")),
 911             vec![los_pollos_suggestion("los")],
 912         );
 913
 914         Ok(())
 915     }
 916
 917     /// Tests re-ingesting suggestions from an updated attachment.
 918     #[test]
 919     fn reingest_amp_suggestions() -> anyhow::Result<()> {
 920         before_each();
 921
 922         let mut store = TestStore::new(MockRemoteSettingsClient::default().with_record(
 923             "data",
 924             "1234",
 925             json!([los_pollos_amp(), good_place_eats_amp()]),
 926         ));
 927         // Ingest once
 928         store.ingest(SuggestIngestionConstraints::default());
 929         // Update the snapshot with new suggestions: Los pollos has a new name and Good place eats
 930         // is now serving Penne
 931         store.replace_client(MockRemoteSettingsClient::default().with_record(
 932             "data",
 933             "1234",
 934             json!([
 935                 los_pollos_amp().merge(json!({
 936                     "title": "Los Pollos Hermanos - Now Serving at 14 Locations!",
 937                 })),
 938                 good_place_eats_amp().merge(json!({
 939                     "keywords": ["pe", "pen", "penne", "penne for your thoughts"],
 940                     "title": "Penne for Your Thoughts",
 941                     "url": "https://penne.biz",
 942                 }))
 943             ]),
 944         ));
 945         store.ingest(SuggestIngestionConstraints::default());
 946
 947         assert!(matches!(
 948             store.fetch_suggestions(SuggestionQuery::amp("lo")).as_slice(),
 949             [ Suggestion::Amp { title, .. } ] if title == "Los Pollos Hermanos - Now Serving at 14 Locations!",
 950         ));
 951
 952         assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("la")), vec![]);
 953         assert!(matches!(
 954             store.fetch_suggestions(SuggestionQuery::amp("pe")).as_slice(),
 955             [ Suggestion::Amp { title, url, .. } ] if title == "Penne for Your Thoughts" && url == "https://penne.biz"
 956         ));
 957
 958         Ok(())
 959     }
 960
 961     /// Tests re-ingesting icons from an updated attachment.
 962     #[test]
 963     fn reingest_icons() -> anyhow::Result<()> {
 964         before_each();
 965
 966         let mut store = TestStore::new(
 967             MockRemoteSettingsClient::default()
 968                 .with_record(
 969                     "data",
 970                     "1234",
 971                     json!([los_pollos_amp(), good_place_eats_amp()]),
 972                 )
 973                 .with_icon(los_pollos_icon())
 974                 .with_icon(good_place_eats_icon()),
 975         );
 976         // This ingestion should run, since the DB is empty
 977         store.ingest(SuggestIngestionConstraints::default());
 978
 979         // Reingest with updated icon data
 980         //  - Los pollos gets new data and a new id
 981         //  - Good place eats gets new data only
 982         store.replace_client(
 983             MockRemoteSettingsClient::default()
 984                 .with_record(
 985                     "data",
 986                     "1234",
 987                     json!([
 988                         los_pollos_amp().merge(json!({"icon": "1000"})),
 989                         good_place_eats_amp()
 990                     ]),
 991                 )
 992                 .with_icon(MockIcon {
 993                     id: "1000",
 994                     data: "new-los-pollos-icon",
 995                     ..los_pollos_icon()
 996                 })
 997                 .with_icon(MockIcon {
 998                     data: "new-good-place-eats-icon",
 999                     ..good_place_eats_icon()
1000                 }),
1001         );
1002         store.ingest(SuggestIngestionConstraints::default());
1003
1004         assert!(matches!(
1005             store.fetch_suggestions(SuggestionQuery::amp("lo")).as_slice(),
1006             [ Suggestion::Amp { icon, .. } ] if *icon == Some("new-los-pollos-icon".as_bytes().to_vec())
1007         ));
1008
1009         assert!(matches!(
1010             store.fetch_suggestions(SuggestionQuery::amp("la")).as_slice(),
1011             [ Suggestion::Amp { icon, .. } ] if *icon == Some("new-good-place-eats-icon".as_bytes().to_vec())
1012         ));
1013
1014         Ok(())
1015     }
1016
1017     /// Tests re-ingesting AMO suggestions from an updated attachment.
1018     #[test]
1019     fn reingest_amo_suggestions() -> anyhow::Result<()> {
1020         before_each();
1021
1022         let mut store = TestStore::new(
1023             MockRemoteSettingsClient::default()
1024                 .with_record("amo-suggestions", "data-1", json!([relay_amo()]))
1025                 .with_record(
1026                     "amo-suggestions",
1027                     "data-2",
1028                     json!([dark_mode_amo(), foxy_guestures_amo()]),
1029                 ),
1030         );
1031
1032         store.ingest(SuggestIngestionConstraints::default());
1033
1034         assert_eq!(
1035             store.fetch_suggestions(SuggestionQuery::amo("masking e")),
1036             vec![relay_suggestion()],
1037         );
1038         assert_eq!(
1039             store.fetch_suggestions(SuggestionQuery::amo("night")),
1040             vec![dark_mode_suggestion()],
1041         );
1042         assert_eq!(
1043             store.fetch_suggestions(SuggestionQuery::amo("grammar")),
1044             vec![foxy_guestures_suggestion()],
1045         );
1046
1047         // Update the snapshot with new suggestions: update the second, drop the
1048         // third, and add the fourth.
1049         store.replace_client(
1050             MockRemoteSettingsClient::default()
1051                 .with_record("amo-suggestions", "data-1", json!([relay_amo()]))
1052                 .with_record(
1053                     "amo-suggestions",
1054                     "data-2",
1055                     json!([
1056                         dark_mode_amo().merge(json!({"title": "Updated second suggestion"})),
1057                         new_tab_override_amo(),
1058                     ]),
1059                 ),
1060         );
1061         store.ingest(SuggestIngestionConstraints::default());
1062
1063         assert_eq!(
1064             store.fetch_suggestions(SuggestionQuery::amo("masking e")),
1065             vec![relay_suggestion()],
1066         );
1067         assert!(matches!(
1068             store.fetch_suggestions(SuggestionQuery::amo("night")).as_slice(),
1069             [Suggestion::Amo { title, .. } ] if title == "Updated second suggestion"
1070         ));
1071         assert_eq!(
1072             store.fetch_suggestions(SuggestionQuery::amo("grammar")),
1073             vec![],
1074         );
1075         assert_eq!(
1076             store.fetch_suggestions(SuggestionQuery::amo("image search")),
1077             vec![new_tab_override_suggestion()],
1078         );
1079
1080         Ok(())
1081     }
1082
1083     /// Tests ingesting tombstones for previously-ingested suggestions and
1084     /// icons.
1085     #[test]
1086     fn ingest_tombstones() -> anyhow::Result<()> {
1087         before_each();
1088
1089         let mut store = TestStore::new(
1090             MockRemoteSettingsClient::default()
1091                 .with_record("data", "data-1", json!([los_pollos_amp()]))
1092                 .with_record("data", "data-2", json!([good_place_eats_amp()]))
1093                 .with_icon(los_pollos_icon())
1094                 .with_icon(good_place_eats_icon()),
1095         );
1096         store.ingest(SuggestIngestionConstraints::default());
1097         assert_eq!(
1098             store.fetch_suggestions(SuggestionQuery::amp("lo")),
1099             vec![los_pollos_suggestion("los")],
1100         );
1101         assert_eq!(
1102             store.fetch_suggestions(SuggestionQuery::amp("la")),
1103             vec![good_place_eats_suggestion("lasagna")],
1104         );
1105         // Re-ingest with:
1106         //   - Los pollos replaced with a tombstone
1107         //   - Good place eat's icon replaced with a tombstone
1108         store.replace_client(
1109             MockRemoteSettingsClient::default()
1110                 .with_tombstone("data", "data-1")
1111                 .with_record("data", "data-2", json!([good_place_eats_amp()]))
1112                 .with_icon_tombstone(los_pollos_icon())
1113                 .with_icon_tombstone(good_place_eats_icon()),
1114         );
1115         store.ingest(SuggestIngestionConstraints::default());
1116
1117         assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("lo")), vec![]);
1118         assert!(matches!(
1119             store.fetch_suggestions(SuggestionQuery::amp("la")).as_slice(),
1120             [
1121                 Suggestion::Amp { icon, icon_mimetype, .. }
1122             ] if icon.is_none() && icon_mimetype.is_none(),
1123         ));
1124         Ok(())
1125     }
1126
1127     /// Tests clearing the store.
1128     #[test]
1129     fn clear() -> anyhow::Result<()> {
1130         before_each();
1131
1132         let store = TestStore::new(
1133             MockRemoteSettingsClient::default()
1134                 .with_record("data", "data-1", json!([los_pollos_amp()]))
1135                 .with_record("data", "data-2", json!([good_place_eats_amp()]))
1136                 .with_icon(los_pollos_icon())
1137                 .with_icon(good_place_eats_icon()),
1138         );
1139         store.ingest(SuggestIngestionConstraints::default());
1140         assert!(store.count_rows("suggestions") > 0);
1141         assert!(store.count_rows("keywords") > 0);
1142         assert!(store.count_rows("icons") > 0);
1143
1144         store.inner.clear()?;
1145         assert!(store.count_rows("suggestions") == 0);
1146         assert!(store.count_rows("keywords") == 0);
1147         assert!(store.count_rows("icons") == 0);
1148
1149         Ok(())
1150     }
1151
1152     /// Tests querying suggestions.
1153     #[test]
1154     fn query() -> anyhow::Result<()> {
1155         before_each();
1156
1157         let store = TestStore::new(
1158             MockRemoteSettingsClient::default()
1159                 .with_record(
1160                     "data",
1161                     "data-1",
1162                     json!([
1163                         good_place_eats_amp(),
1164                         california_wiki(),
1165                         caltech_wiki(),
1166                         multimatch_wiki(),
1167                     ]),
1168                 )
1169                 .with_record(
1170                     "amo-suggestions",
1171                     "data-2",
1172                     json!([relay_amo(), multimatch_amo(),]),
1173                 )
1174                 .with_record(
1175                     "pocket-suggestions",
1176                     "data-3",
1177                     json!([burnout_pocket(), multimatch_pocket(),]),
1178                 )
1179                 .with_record("yelp-suggestions", "data-4", json!([ramen_yelp(),]))
1180                 .with_record("yeld-suggestions", "data-4", json!([ramen_yelp(),]))
1181                 .with_record("mdn-suggestions", "data-5", json!([array_mdn(),]))
1182                 .with_icon(good_place_eats_icon())
1183                 .with_icon(california_icon())
1184                 .with_icon(caltech_icon())
1185                 .with_icon(yelp_favicon())
1186                 .with_icon(multimatch_wiki_icon()),
1187         );
1188
1189         store.ingest(SuggestIngestionConstraints::default());
1190
1191         assert_eq!(
1192             store.fetch_suggestions(SuggestionQuery::all_providers("")),
1193             vec![]
1194         );
1195         assert_eq!(
1196             store.fetch_suggestions(SuggestionQuery::all_providers("la")),
1197             vec![good_place_eats_suggestion("lasagna"),]
1198         );
1199         assert_eq!(
1200             store.fetch_suggestions(SuggestionQuery::all_providers("multimatch")),
1201             vec![
1202                 multimatch_pocket_suggestion(true),
1203                 multimatch_amo_suggestion(),
1204                 multimatch_wiki_suggestion(),
1205             ]
1206         );
1207         assert_eq!(
1208             store.fetch_suggestions(SuggestionQuery::all_providers("MultiMatch")),
1209             vec![
1210                 multimatch_pocket_suggestion(true),
1211                 multimatch_amo_suggestion(),
1212                 multimatch_wiki_suggestion(),
1213             ]
1214         );
1215         assert_eq!(
1216             store.fetch_suggestions(SuggestionQuery::all_providers("multimatch").limit(2)),
1217             vec![
1218                 multimatch_pocket_suggestion(true),
1219                 multimatch_amo_suggestion(),
1220             ],
1221         );
1222         assert_eq!(
1223             store.fetch_suggestions(SuggestionQuery::amp("la")),
1224             vec![good_place_eats_suggestion("lasagna")],
1225         );
1226         assert_eq!(
1227             store.fetch_suggestions(SuggestionQuery::all_providers_except(
1228                 "la",
1229                 SuggestionProvider::Amp
1230             )),
1231             vec![],
1232         );
1233         assert_eq!(
1234             store.fetch_suggestions(SuggestionQuery::with_providers("la", vec![])),
1235             vec![],
1236         );
1237         assert_eq!(
1238             store.fetch_suggestions(SuggestionQuery::with_providers(
1239                 "cal",
1240                 vec![
1241                     SuggestionProvider::Amp,
1242                     SuggestionProvider::Amo,
1243                     SuggestionProvider::Pocket,
1244                 ]
1245             )),
1246             vec![],
1247         );
1248         assert_eq!(
1249             store.fetch_suggestions(SuggestionQuery::wikipedia("cal")),
1250             vec![
1251                 california_suggestion("california"),
1252                 caltech_suggestion("california"),
1253             ],
1254         );
1255         assert_eq!(
1256             store.fetch_suggestions(SuggestionQuery::wikipedia("cal").limit(1)),
1257             vec![california_suggestion("california"),],
1258         );
1259         assert_eq!(
1260             store.fetch_suggestions(SuggestionQuery::with_providers("cal", vec![])),
1261             vec![],
1262         );
1263         assert_eq!(
1264             store.fetch_suggestions(SuggestionQuery::amo("spam")),
1265             vec![relay_suggestion()],
1266         );
1267         assert_eq!(
1268             store.fetch_suggestions(SuggestionQuery::amo("masking")),
1269             vec![relay_suggestion()],
1270         );
1271         assert_eq!(
1272             store.fetch_suggestions(SuggestionQuery::amo("masking e")),
1273             vec![relay_suggestion()],
1274         );
1275         assert_eq!(
1276             store.fetch_suggestions(SuggestionQuery::amo("masking s")),
1277             vec![],
1278         );
1279         assert_eq!(
1280             store.fetch_suggestions(SuggestionQuery::with_providers(
1281                 "soft",
1282                 vec![SuggestionProvider::Amp, SuggestionProvider::Wikipedia]
1283             )),
1284             vec![],
1285         );
1286         assert_eq!(
1287             store.fetch_suggestions(SuggestionQuery::pocket("soft")),
1288             vec![burnout_suggestion(false),],
1289         );
1290         assert_eq!(
1291             store.fetch_suggestions(SuggestionQuery::pocket("soft l")),
1292             vec![burnout_suggestion(false),],
1293         );
1294         assert_eq!(
1295             store.fetch_suggestions(SuggestionQuery::pocket("sof")),
1296             vec![],
1297         );
1298         assert_eq!(
1299             store.fetch_suggestions(SuggestionQuery::pocket("burnout women")),
1300             vec![burnout_suggestion(true),],
1301         );
1302         assert_eq!(
1303             store.fetch_suggestions(SuggestionQuery::pocket("burnout person")),
1304             vec![],
1305         );
1306         assert_eq!(
1307             store.fetch_suggestions(SuggestionQuery::yelp("best spicy ramen delivery in tokyo")),
1308             vec![ramen_suggestion(
1309                 "best spicy ramen delivery in tokyo",
1310                 "https://www.yelp.com/search?find_desc=best+spicy+ramen+delivery&find_loc=tokyo"
1311             ),],
1312         );
1313         assert_eq!(
1314             store.fetch_suggestions(SuggestionQuery::yelp("BeSt SpIcY rAmEn DeLiVeRy In ToKyO")),
1315             vec![ramen_suggestion(
1316                 "BeSt SpIcY rAmEn DeLiVeRy In ToKyO",
1317                 "https://www.yelp.com/search?find_desc=BeSt+SpIcY+rAmEn+DeLiVeRy&find_loc=ToKyO"
1318             ),],
1319         );
1320         assert_eq!(
1321             store.fetch_suggestions(SuggestionQuery::yelp("best ramen delivery in tokyo")),
1322             vec![ramen_suggestion(
1323                 "best ramen delivery in tokyo",
1324                 "https://www.yelp.com/search?find_desc=best+ramen+delivery&find_loc=tokyo"
1325             ),],
1326         );
1327         assert_eq!(
1328             store.fetch_suggestions(SuggestionQuery::yelp(
1329                 "best invalid_ramen delivery in tokyo"
1330             )),
1331             vec![],
1332         );
1333         assert_eq!(
1334             store.fetch_suggestions(SuggestionQuery::yelp("best in tokyo")),
1335             vec![],
1336         );
1337         assert_eq!(
1338             store.fetch_suggestions(SuggestionQuery::yelp("super best ramen in tokyo")),
1339             vec![ramen_suggestion(
1340                 "super best ramen in tokyo",
1341                 "https://www.yelp.com/search?find_desc=super+best+ramen&find_loc=tokyo"
1342             ),],
1343         );
1344         assert_eq!(
1345             store.fetch_suggestions(SuggestionQuery::yelp("invalid_best ramen in tokyo")),
1346             vec![],
1347         );
1348         assert_eq!(
1349             store.fetch_suggestions(SuggestionQuery::yelp("ramen delivery in tokyo")),
1350             vec![ramen_suggestion(
1351                 "ramen delivery in tokyo",
1352                 "https://www.yelp.com/search?find_desc=ramen+delivery&find_loc=tokyo"
1353             ),],
1354         );
1355         assert_eq!(
1356             store.fetch_suggestions(SuggestionQuery::yelp("ramen super delivery in tokyo")),
1357             vec![ramen_suggestion(
1358                 "ramen super delivery in tokyo",
1359                 "https://www.yelp.com/search?find_desc=ramen+super+delivery&find_loc=tokyo"
1360             ),],
1361         );
1362         assert_eq!(
1363             store.fetch_suggestions(SuggestionQuery::yelp("ramen invalid_delivery in tokyo")),
1364             vec![],
1365         );
1366         assert_eq!(
1367             store.fetch_suggestions(SuggestionQuery::yelp("ramen in tokyo")),
1368             vec![ramen_suggestion(
1369                 "ramen in tokyo",
1370                 "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1371             ),],
1372         );
1373         assert_eq!(
1374             store.fetch_suggestions(SuggestionQuery::yelp("ramen near tokyo")),
1375             vec![ramen_suggestion(
1376                 "ramen near tokyo",
1377                 "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1378             ),],
1379         );
1380         assert_eq!(
1381             store.fetch_suggestions(SuggestionQuery::yelp("ramen invalid_in tokyo")),
1382             vec![],
1383         );
1384         assert_eq!(
1385             store.fetch_suggestions(SuggestionQuery::yelp("ramen in San Francisco")),
1386             vec![ramen_suggestion(
1387                 "ramen in San Francisco",
1388                 "https://www.yelp.com/search?find_desc=ramen&find_loc=San+Francisco"
1389             ),],
1390         );
1391         assert_eq!(
1392             store.fetch_suggestions(SuggestionQuery::yelp("ramen in")),
1393             vec![ramen_suggestion(
1394                 "ramen in",
1395                 "https://www.yelp.com/search?find_desc=ramen"
1396             ),],
1397         );
1398         assert_eq!(
1399             store.fetch_suggestions(SuggestionQuery::yelp("ramen near by")),
1400             vec![ramen_suggestion(
1401                 "ramen near by",
1402                 "https://www.yelp.com/search?find_desc=ramen+near+by"
1403             )
1404             .has_location_sign(false),],
1405         );
1406         assert_eq!(
1407             store.fetch_suggestions(SuggestionQuery::yelp("ramen near me")),
1408             vec![ramen_suggestion(
1409                 "ramen near me",
1410                 "https://www.yelp.com/search?find_desc=ramen+near+me"
1411             )
1412             .has_location_sign(false),],
1413         );
1414         assert_eq!(
1415             store.fetch_suggestions(SuggestionQuery::yelp("ramen near by tokyo")),
1416             vec![],
1417         );
1418         assert_eq!(
1419             store.fetch_suggestions(SuggestionQuery::yelp("ramen")),
1420             vec![
1421                 ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
1422                     .has_location_sign(false),
1423             ],
1424         );
1425         // Test an extremely long yelp query
1426         assert_eq!(
1427             store.fetch_suggestions(SuggestionQuery::yelp(
1428                 "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789"
1429             )),
1430             vec![
1431                 ramen_suggestion(
1432                     "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789",
1433                     "https://www.yelp.com/search?find_desc=012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789"
1434                 ).has_location_sign(false),
1435             ],
1436         );
1437         // This query is over the limit and no suggestions should be returned
1438         assert_eq!(
1439             store.fetch_suggestions(SuggestionQuery::yelp(
1440                 "012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789Z"
1441             )),
1442             vec![],
1443         );
1444         assert_eq!(
1445             store.fetch_suggestions(SuggestionQuery::yelp("best delivery")),
1446             vec![],
1447         );
1448         assert_eq!(
1449             store.fetch_suggestions(SuggestionQuery::yelp("same_modifier same_modifier")),
1450             vec![],
1451         );
1452         assert_eq!(
1453             store.fetch_suggestions(SuggestionQuery::yelp("same_modifier ")),
1454             vec![],
1455         );
1456         assert_eq!(
1457             store.fetch_suggestions(SuggestionQuery::yelp("yelp ramen")),
1458             vec![
1459                 ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
1460                     .has_location_sign(false),
1461             ],
1462         );
1463         assert_eq!(
1464             store.fetch_suggestions(SuggestionQuery::yelp("yelp keyword ramen")),
1465             vec![
1466                 ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
1467                     .has_location_sign(false),
1468             ],
1469         );
1470         assert_eq!(
1471             store.fetch_suggestions(SuggestionQuery::yelp("ramen in tokyo yelp")),
1472             vec![ramen_suggestion(
1473                 "ramen in tokyo",
1474                 "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1475             )],
1476         );
1477         assert_eq!(
1478             store.fetch_suggestions(SuggestionQuery::yelp("ramen in tokyo yelp keyword")),
1479             vec![ramen_suggestion(
1480                 "ramen in tokyo",
1481                 "https://www.yelp.com/search?find_desc=ramen&find_loc=tokyo"
1482             )],
1483         );
1484         assert_eq!(
1485             store.fetch_suggestions(SuggestionQuery::yelp("yelp ramen yelp")),
1486             vec![
1487                 ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
1488                     .has_location_sign(false)
1489             ],
1490         );
1491         assert_eq!(
1492             store.fetch_suggestions(SuggestionQuery::yelp("best yelp ramen")),
1493             vec![],
1494         );
1495         assert_eq!(
1496             store.fetch_suggestions(SuggestionQuery::yelp("Spicy R")),
1497             vec![ramen_suggestion(
1498                 "Spicy Ramen",
1499                 "https://www.yelp.com/search?find_desc=Spicy+Ramen"
1500             )
1501             .has_location_sign(false)
1502             .subject_exact_match(false)],
1503         );
1504         assert_eq!(
1505             store.fetch_suggestions(SuggestionQuery::yelp("BeSt             Ramen")),
1506             vec![ramen_suggestion(
1507                 "BeSt Ramen",
1508                 "https://www.yelp.com/search?find_desc=BeSt+Ramen"
1509             )
1510             .has_location_sign(false)],
1511         );
1512         assert_eq!(
1513             store.fetch_suggestions(SuggestionQuery::yelp("BeSt             Spicy R")),
1514             vec![ramen_suggestion(
1515                 "BeSt Spicy Ramen",
1516                 "https://www.yelp.com/search?find_desc=BeSt+Spicy+Ramen"
1517             )
1518             .has_location_sign(false)
1519             .subject_exact_match(false)],
1520         );
1521         assert_eq!(
1522             store.fetch_suggestions(SuggestionQuery::yelp("BeSt             R")),
1523             vec![],
1524         );
1525         assert_eq!(store.fetch_suggestions(SuggestionQuery::yelp("r")), vec![],);
1526         assert_eq!(
1527             store.fetch_suggestions(SuggestionQuery::yelp("ra")),
1528             vec![
1529                 ramen_suggestion("rats", "https://www.yelp.com/search?find_desc=rats")
1530                     .has_location_sign(false)
1531                     .subject_exact_match(false)
1532             ],
1533         );
1534         assert_eq!(
1535             store.fetch_suggestions(SuggestionQuery::yelp("ram")),
1536             vec![
1537                 ramen_suggestion("ramen", "https://www.yelp.com/search?find_desc=ramen")
1538                     .has_location_sign(false)
1539                     .subject_exact_match(false)
1540             ],
1541         );
1542         assert_eq!(
1543             store.fetch_suggestions(SuggestionQuery::yelp("rac")),
1544             vec![
1545                 ramen_suggestion("raccoon", "https://www.yelp.com/search?find_desc=raccoon")
1546                     .has_location_sign(false)
1547                     .subject_exact_match(false)
1548             ],
1549         );
1550         assert_eq!(
1551             store.fetch_suggestions(SuggestionQuery::yelp("best r")),
1552             vec![],
1553         );
1554         assert_eq!(
1555             store.fetch_suggestions(SuggestionQuery::yelp("best ra")),
1556             vec![ramen_suggestion(
1557                 "best rats",
1558                 "https://www.yelp.com/search?find_desc=best+rats"
1559             )
1560             .has_location_sign(false)
1561             .subject_exact_match(false)],
1562         );
1563
1564         Ok(())
1565     }
1566
1567     // Tests querying AMP / Wikipedia / Pocket
1568     #[test]
1569     fn query_with_multiple_providers_and_diff_scores() -> anyhow::Result<()> {
1570         before_each();
1571
1572         let store = TestStore::new(
1573             // Create a data set where one keyword matches multiple suggestions from each provider
1574             // where the scores are manually set.  We will test that the fetched suggestions are in
1575             // the correct order.
1576             MockRemoteSettingsClient::default()
1577                 .with_record(
1578                     "data",
1579                     "data-1",
1580                     json!([
1581                         los_pollos_amp().merge(json!({
1582                             "keywords": ["amp wiki match"],
1583                             "score": 0.3,
1584                         })),
1585                         good_place_eats_amp().merge(json!({
1586                             "keywords": ["amp wiki match"],
1587                             "score": 0.1,
1588                         })),
1589                         california_wiki().merge(json!({
1590                             "keywords": ["amp wiki match", "pocket wiki match"],
1591                         })),
1592                     ]),
1593                 )
1594                 .with_record(
1595                     "pocket-suggestions",
1596                     "data-3",
1597                     json!([
1598                         burnout_pocket().merge(json!({
1599                             "lowConfidenceKeywords": ["work-life balance", "pocket wiki match"],
1600                             "score": 0.05,
1601                         })),
1602                         multimatch_pocket().merge(json!({
1603                             "highConfidenceKeywords": ["pocket wiki match"],
1604                             "score": 0.88,
1605                         })),
1606                     ]),
1607                 )
1608                 .with_icon(los_pollos_icon())
1609                 .with_icon(good_place_eats_icon())
1610                 .with_icon(california_icon()),
1611         );
1612
1613         store.ingest(SuggestIngestionConstraints::default());
1614         assert_eq!(
1615             store.fetch_suggestions(SuggestionQuery::all_providers("amp wiki match")),
1616             vec![
1617                 los_pollos_suggestion("amp wiki match").with_score(0.3),
1618                 // Wikipedia entries default to a 0.2 score
1619                 california_suggestion("amp wiki match"),
1620                 good_place_eats_suggestion("amp wiki match").with_score(0.1),
1621             ]
1622         );
1623         assert_eq!(
1624             store.fetch_suggestions(SuggestionQuery::all_providers("amp wiki match").limit(2)),
1625             vec![
1626                 los_pollos_suggestion("amp wiki match").with_score(0.3),
1627                 california_suggestion("amp wiki match"),
1628             ]
1629         );
1630         assert_eq!(
1631             store.fetch_suggestions(SuggestionQuery::all_providers("pocket wiki match")),
1632             vec![
1633                 multimatch_pocket_suggestion(true).with_score(0.88),
1634                 california_suggestion("pocket wiki match"),
1635                 burnout_suggestion(false).with_score(0.05),
1636             ]
1637         );
1638         assert_eq!(
1639             store.fetch_suggestions(SuggestionQuery::all_providers("pocket wiki match").limit(1)),
1640             vec![multimatch_pocket_suggestion(true).with_score(0.88),]
1641         );
1642         // test duplicate providers
1643         assert_eq!(
1644             store.fetch_suggestions(SuggestionQuery::with_providers(
1645                 "work-life balance",
1646                 vec![SuggestionProvider::Pocket, SuggestionProvider::Pocket],
1647             )),
1648             vec![burnout_suggestion(false).with_score(0.05),]
1649         );
1650
1651         Ok(())
1652     }
1653
1654     // Tests querying multiple suggestions with multiple keywords with same prefix keyword
1655     #[test]
1656     fn query_with_amp_mobile_provider() -> anyhow::Result<()> {
1657         before_each();
1658
1659         // Use the exact same data for both the Amp and AmpMobile record
1660         let store = TestStore::new(
1661             MockRemoteSettingsClient::default()
1662                 .with_record(
1663                     "amp-mobile-suggestions",
1664                     "amp-mobile-1",
1665                     json!([good_place_eats_amp()]),
1666                 )
1667                 .with_record("data", "data-1", json!([good_place_eats_amp()]))
1668                 // This icon is shared by both records which is kind of weird and probably not how
1669                 // things would work in practice, but it's okay for the tests.
1670                 .with_icon(good_place_eats_icon()),
1671         );
1672         store.ingest(SuggestIngestionConstraints::default());
1673         // The query results should be exactly the same for both the Amp and AmpMobile data
1674         assert_eq!(
1675             store.fetch_suggestions(SuggestionQuery::amp_mobile("las")),
1676             vec![good_place_eats_suggestion("lasagna")]
1677         );
1678         assert_eq!(
1679             store.fetch_suggestions(SuggestionQuery::amp("las")),
1680             vec![good_place_eats_suggestion("lasagna")]
1681         );
1682         Ok(())
1683     }
1684
1685     /// Tests ingesting malformed Remote Settings records that we understand,
1686     /// but that are missing fields, or aren't in the format we expect.
1687     #[test]
1688     fn ingest_malformed() -> anyhow::Result<()> {
1689         before_each();
1690
1691         let store = TestStore::new(
1692             MockRemoteSettingsClient::default()
1693                 // Amp/Wikipedia record without an attachment.
1694                 .with_record_but_no_attachment("data", "data-1")
1695                 // Icon record without an attachment.
1696                 .with_record_but_no_attachment("icon", "icon-1")
1697                 // Icon record with an ID that's not `icon-{id}`, so suggestions in
1698                 // the data attachment won't be able to reference it.
1699                 .with_record("icon", "bad-icon-id", json!("i-am-an-icon")),
1700         );
1701
1702         store.ingest(SuggestIngestionConstraints::default());
1703
1704         store.read(|dao| {
1705             assert_eq!(
1706                 dao.get_meta::<u64>(SuggestRecordType::Icon.last_ingest_meta_key().as_str())?,
1707                 Some(store.last_modified_timestamp())
1708             );
1709             assert_eq!(
1710                 dao.conn
1711                     .query_one::<i64>("SELECT count(*) FROM suggestions")?,
1712                 0
1713             );
1714             assert_eq!(dao.conn.query_one::<i64>("SELECT count(*) FROM icons")?, 0);
1715
1716             Ok(())
1717         })?;
1718
1719         Ok(())
1720     }
1721
1722     /// Tests that we only ingest providers that we're concerned with.
1723     #[test]
1724     fn ingest_constraints_provider() -> anyhow::Result<()> {
1725         before_each();
1726
1727         let store = TestStore::new(
1728             MockRemoteSettingsClient::default()
1729                 .with_record("data", "data-1", json!([los_pollos_amp()]))
1730                 .with_record("yelp", "yelp-1", json!([ramen_yelp()]))
1731                 .with_icon(los_pollos_icon()),
1732         );
1733
1734         // Write a last ingestion times to test that we overwrite it properly
1735         store.write(|dao| {
1736             // Check that existing data is updated properly.
1737             dao.put_meta(
1738                 SuggestRecordType::AmpWikipedia
1739                     .last_ingest_meta_key()
1740                     .as_str(),
1741                 1,
1742             )?;
1743             Ok(())
1744         })?;
1745
1746         let constraints = SuggestIngestionConstraints {
1747             max_suggestions: Some(100),
1748             providers: Some(vec![SuggestionProvider::Amp, SuggestionProvider::Pocket]),
1749             ..SuggestIngestionConstraints::default()
1750         };
1751         store.ingest(constraints);
1752
1753         // This should have been ingested
1754         assert_eq!(
1755             store.fetch_suggestions(SuggestionQuery::amp("lo")),
1756             vec![los_pollos_suggestion("los")]
1757         );
1758         // This should not have been ingested, since it wasn't in the providers list
1759         assert_eq!(
1760             store.fetch_suggestions(SuggestionQuery::yelp("best ramen")),
1761             vec![]
1762         );
1763
1764         store.read(|dao| {
1765             // This should have its last_modified_timestamp updated, since we ingested an amp
1766             // record
1767             assert_eq!(
1768                 dao.get_meta::<u64>(
1769                     SuggestRecordType::AmpWikipedia
1770                         .last_ingest_meta_key()
1771                         .as_str()
1772                 )?,
1773                 Some(store.last_modified_timestamp())
1774             );
1775             // This should have its last_modified_timestamp updated, since we ingested an icon
1776             assert_eq!(
1777                 dao.get_meta::<u64>(SuggestRecordType::Icon.last_ingest_meta_key().as_str())?,
1778                 Some(store.last_modified_timestamp())
1779             );
1780             // This should not have its last_modified_timestamp updated, since there were no pocket
1781             // items to ingest
1782             assert_eq!(
1783                 dao.get_meta::<u64>(SuggestRecordType::Pocket.last_ingest_meta_key().as_str())?,
1784                 None
1785             );
1786             // This should not have its last_modified_timestamp updated, since we did not ask to
1787             // ingest yelp items.
1788             assert_eq!(
1789                 dao.get_meta::<u64>(SuggestRecordType::Yelp.last_ingest_meta_key().as_str())?,
1790                 None
1791             );
1792             assert_eq!(
1793                 dao.get_meta::<u64>(SuggestRecordType::Amo.last_ingest_meta_key().as_str())?,
1794                 None
1795             );
1796             assert_eq!(
1797                 dao.get_meta::<u64>(SuggestRecordType::Mdn.last_ingest_meta_key().as_str())?,
1798                 None
1799             );
1800             assert_eq!(
1801                 dao.get_meta::<u64>(SuggestRecordType::AmpMobile.last_ingest_meta_key().as_str())?,
1802                 None
1803             );
1804             assert_eq!(
1805                 dao.get_meta::<u64>(
1806                     SuggestRecordType::GlobalConfig
1807                         .last_ingest_meta_key()
1808                         .as_str()
1809                 )?,
1810                 None
1811             );
1812             Ok(())
1813         })?;
1814
1815         Ok(())
1816     }
1817
1818     /// Tests that records with invalid attachments are ignored
1819     #[test]
1820     fn skip_over_invalid_records() -> anyhow::Result<()> {
1821         before_each();
1822
1823         let store = TestStore::new(
1824             MockRemoteSettingsClient::default()
1825                 // valid record
1826                 .with_record("data", "data-1", json!([good_place_eats_amp()]))
1827                 // This attachment is missing the `title` field and is invalid
1828                 .with_record(
1829                     "data",
1830                     "data-2",
1831                     json!([{
1832                             "id": 1,
1833                             "advertiser": "Los Pollos Hermanos",
1834                             "iab_category": "8 - Food & Drink",
1835                             "keywords": ["lo", "los", "los pollos"],
1836                             "url": "https://www.lph-nm.biz",
1837                             "icon": "5678",
1838                             "impression_url": "https://example.com/impression_url",
1839                             "click_url": "https://example.com/click_url",
1840                             "score": 0.3
1841                     }]),
1842                 )
1843                 .with_icon(good_place_eats_icon()),
1844         );
1845
1846         store.ingest(SuggestIngestionConstraints::default());
1847
1848         // Test that the valid record was read
1849         assert_eq!(
1850             store.fetch_suggestions(SuggestionQuery::amp("la")),
1851             vec![good_place_eats_suggestion("lasagna")]
1852         );
1853         // Test that the invalid record was skipped
1854         assert_eq!(store.fetch_suggestions(SuggestionQuery::amp("lo")), vec![]);
1855
1856         Ok(())
1857     }
1858
1859     #[test]
1860     fn query_mdn() -> anyhow::Result<()> {
1861         before_each();
1862
1863         let store = TestStore::new(MockRemoteSettingsClient::default().with_record(
1864             "mdn-suggestions",
1865             "mdn-1",
1866             json!([array_mdn()]),
1867         ));
1868         store.ingest(SuggestIngestionConstraints::default());
1869         // prefix
1870         assert_eq!(
1871             store.fetch_suggestions(SuggestionQuery::mdn("array")),
1872             vec![array_suggestion(),]
1873         );
1874         // prefix + partial suffix
1875         assert_eq!(
1876             store.fetch_suggestions(SuggestionQuery::mdn("array java")),
1877             vec![array_suggestion(),]
1878         );
1879         // prefix + entire suffix
1880         assert_eq!(
1881             store.fetch_suggestions(SuggestionQuery::mdn("javascript array")),
1882             vec![array_suggestion(),]
1883         );
1884         // partial prefix word
1885         assert_eq!(
1886             store.fetch_suggestions(SuggestionQuery::mdn("wild")),
1887             vec![]
1888         );
1889         // single word
1890         assert_eq!(
1891             store.fetch_suggestions(SuggestionQuery::mdn("wildcard")),
1892             vec![array_suggestion()]
1893         );
1894         Ok(())
1895     }
1896
1897     #[test]
1898     fn query_no_yelp_icon_data() -> anyhow::Result<()> {
1899         before_each();
1900
1901         let store = TestStore::new(
1902             MockRemoteSettingsClient::default().with_record(
1903                 "yelp-suggestions",
1904                 "yelp-1",
1905                 json!([ramen_yelp()]),
1906             ), // Note: yelp_favicon() is missing
1907         );
1908         store.ingest(SuggestIngestionConstraints::default());
1909         assert!(matches!(
1910             store.fetch_suggestions(SuggestionQuery::yelp("ramen")).as_slice(),
1911             [Suggestion::Yelp { icon, icon_mimetype, .. }] if icon.is_none() && icon_mimetype.is_none()
1912         ));
1913
1914         Ok(())
1915     }
1916
1917     #[test]
1918     fn weather() -> anyhow::Result<()> {
1919         before_each();
1920
1921         let store = TestStore::new(MockRemoteSettingsClient::default().with_inline_record(
1922             "weather",
1923             "weather-1",
1924             json!({
1925                 "min_keyword_length": 3,
1926                 "keywords": ["ab", "xyz", "weather"],
1927                 "score": "0.24"
1928             }),
1929         ));
1930         store.ingest(SuggestIngestionConstraints::default());
1931         // No match since the query doesn't match any keyword
1932         assert_eq!(
1933             store.fetch_suggestions(SuggestionQuery::weather("xab")),
1934             vec![]
1935         );
1936         assert_eq!(
1937             store.fetch_suggestions(SuggestionQuery::weather("abx")),
1938             vec![]
1939         );
1940         assert_eq!(
1941             store.fetch_suggestions(SuggestionQuery::weather("xxyz")),
1942             vec![]
1943         );
1944         assert_eq!(
1945             store.fetch_suggestions(SuggestionQuery::weather("xyzx")),
1946             vec![]
1947         );
1948         assert_eq!(
1949             store.fetch_suggestions(SuggestionQuery::weather("weatherx")),
1950             vec![]
1951         );
1952         assert_eq!(
1953             store.fetch_suggestions(SuggestionQuery::weather("xweather")),
1954             vec![]
1955         );
1956         assert_eq!(
1957             store.fetch_suggestions(SuggestionQuery::weather("xwea")),
1958             vec![]
1959         );
1960         assert_eq!(
1961             store.fetch_suggestions(SuggestionQuery::weather("x   weather")),
1962             vec![]
1963         );
1964         assert_eq!(
1965             store.fetch_suggestions(SuggestionQuery::weather("   weather x")),
1966             vec![]
1967         );
1968         // No match since the query is too short
1969         assert_eq!(
1970             store.fetch_suggestions(SuggestionQuery::weather("xy")),
1971             vec![]
1972         );
1973         assert_eq!(
1974             store.fetch_suggestions(SuggestionQuery::weather("ab")),
1975             vec![]
1976         );
1977         assert_eq!(
1978             store.fetch_suggestions(SuggestionQuery::weather("we")),
1979             vec![]
1980         );
1981         // Matches
1982         assert_eq!(
1983             store.fetch_suggestions(SuggestionQuery::weather("xyz")),
1984             vec![Suggestion::Weather { score: 0.24 },]
1985         );
1986         assert_eq!(
1987             store.fetch_suggestions(SuggestionQuery::weather("wea")),
1988             vec![Suggestion::Weather { score: 0.24 },]
1989         );
1990         assert_eq!(
1991             store.fetch_suggestions(SuggestionQuery::weather("weat")),
1992             vec![Suggestion::Weather { score: 0.24 },]
1993         );
1994         assert_eq!(
1995             store.fetch_suggestions(SuggestionQuery::weather("weath")),
1996             vec![Suggestion::Weather { score: 0.24 },]
1997         );
1998         assert_eq!(
1999             store.fetch_suggestions(SuggestionQuery::weather("weathe")),
2000             vec![Suggestion::Weather { score: 0.24 },]
2001         );
2002         assert_eq!(
2003             store.fetch_suggestions(SuggestionQuery::weather("weather")),
2004             vec![Suggestion::Weather { score: 0.24 },]
2005         );
2006         assert_eq!(
2007             store.fetch_suggestions(SuggestionQuery::weather("  weather  ")),
2008             vec![Suggestion::Weather { score: 0.24 },]
2009         );
2010
2011         assert_eq!(
2012             store.fetch_provider_config(SuggestionProvider::Weather),
2013             Some(SuggestProviderConfig::Weather {
2014                 min_keyword_length: 3,
2015             })
2016         );
2017
2018         Ok(())
2019     }
2020
2021     #[test]
2022     fn fetch_global_config() -> anyhow::Result<()> {
2023         before_each();
2024
2025         let store = TestStore::new(MockRemoteSettingsClient::default().with_inline_record(
2026             "configuration",
2027             "configuration-1",
2028             json!({
2029                 "show_less_frequently_cap": 3,
2030             }),
2031         ));
2032         store.ingest(SuggestIngestionConstraints::default());
2033         assert_eq!(
2034             store.fetch_global_config(),
2035             SuggestGlobalConfig {
2036                 show_less_frequently_cap: 3,
2037             }
2038         );
2039
2040         Ok(())
2041     }
2042
2043     #[test]
2044     fn fetch_global_config_default() -> anyhow::Result<()> {
2045         before_each();
2046
2047         let store = TestStore::new(MockRemoteSettingsClient::default());
2048         store.ingest(SuggestIngestionConstraints::default());
2049         assert_eq!(
2050             store.fetch_global_config(),
2051             SuggestGlobalConfig {
2052                 show_less_frequently_cap: 0,
2053             }
2054         );
2055
2056         Ok(())
2057     }
2058
2059     #[test]
2060     fn fetch_provider_config_none() -> anyhow::Result<()> {
2061         before_each();
2062
2063         let store = TestStore::new(MockRemoteSettingsClient::default());
2064         store.ingest(SuggestIngestionConstraints::default());
2065         assert_eq!(store.fetch_provider_config(SuggestionProvider::Amp), None);
2066         assert_eq!(
2067             store.fetch_provider_config(SuggestionProvider::Weather),
2068             None
2069         );
2070
2071         Ok(())
2072     }
2073
2074     #[test]
2075     fn fetch_provider_config_other() -> anyhow::Result<()> {
2076         before_each();
2077
2078         let store = TestStore::new(MockRemoteSettingsClient::default().with_inline_record(
2079             "weather",
2080             "weather-1",
2081             json!({
2082                 "min_keyword_length": 3,
2083                 "keywords": ["weather"],
2084                 "score": "0.24"
2085             }),
2086         ));
2087         store.ingest(SuggestIngestionConstraints::default());
2088         // Getting the config for a different provider should return None.
2089         assert_eq!(store.fetch_provider_config(SuggestionProvider::Amp), None);
2090         Ok(())
2091     }
2092
2093     #[test]
2094     fn remove_dismissed_suggestions() -> anyhow::Result<()> {
2095         before_each();
2096
2097         let store = TestStore::new(
2098             MockRemoteSettingsClient::default()
2099                 .with_record(
2100                     "data",
2101                     "data-1",
2102                     json!([
2103                         good_place_eats_amp().merge(json!({"keywords": ["cats"]})),
2104                         california_wiki().merge(json!({"keywords": ["cats"]})),
2105                     ]),
2106                 )
2107                 .with_record(
2108                     "amo-suggestions",
2109                     "amo-1",
2110                     json!([relay_amo().merge(json!({"keywords": ["cats"]})),]),
2111                 )
2112                 .with_record(
2113                     "pocket-suggestions",
2114                     "pocket-1",
2115                     json!([burnout_pocket().merge(json!({
2116                         "lowConfidenceKeywords": ["cats"],
2117                     }))]),
2118                 )
2119                 .with_record(
2120                     "mdn-suggestions",
2121                     "mdn-1",
2122                     json!([array_mdn().merge(json!({"keywords": ["cats"]})),]),
2123                 )
2124                 .with_record(
2125                     "amp-mobile-suggestions",
2126                     "amp-mobile-1",
2127                     json!([a1a_amp_mobile().merge(json!({"keywords": ["cats"]})),]),
2128                 )
2129                 .with_icon(good_place_eats_icon())
2130                 .with_icon(caltech_icon()),
2131         );
2132         store.ingest(SuggestIngestionConstraints::default());
2133
2134         // A query for cats should return all suggestions
2135         let query = SuggestionQuery::all_providers("cats");
2136         let results = store.fetch_suggestions(query.clone());
2137         assert_eq!(results.len(), 6);
2138
2139         for result in results {
2140             store
2141                 .inner
2142                 .dismiss_suggestion(result.raw_url().unwrap().to_string())?;
2143         }
2144
2145         // After dismissing the suggestions, the next query shouldn't return them
2146         assert_eq!(store.fetch_suggestions(query.clone()).len(), 0);
2147
2148         // Clearing the dismissals should cause them to be returned again
2149         store.inner.clear_dismissed_suggestions()?;
2150         assert_eq!(store.fetch_suggestions(query.clone()).len(), 6);
2151
2152         Ok(())
2153     }
2154 }