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