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