diff --git a/nidx/nidx_relation/tests/common/mod.rs b/nidx/nidx_relation/tests/common/mod.rs index 779a64fe6d..3df20738c7 100644 --- a/nidx/nidx_relation/tests/common/mod.rs +++ b/nidx/nidx_relation/tests/common/mod.rs @@ -49,3 +49,19 @@ pub fn create_relation( }), } } + +pub fn create_relation_with_metadata( + source: String, + source_node_type: NodeType, + source_subtype: String, + to: String, + to_node_type: NodeType, + to_subtype: String, + rel_type: RelationType, + metadata: RelationMetadata, +) -> Relation { + let mut relation = + create_relation(source, source_node_type, source_subtype, to, to_node_type, to_subtype, rel_type); + relation.metadata = Some(metadata); + relation +} diff --git a/nucliadb_relations2/tests/common/mod.rs b/nucliadb_relations2/tests/common/mod.rs index 4edbca3df5..47a4f1d898 100644 --- a/nucliadb_relations2/tests/common/mod.rs +++ b/nucliadb_relations2/tests/common/mod.rs @@ -49,3 +49,19 @@ pub fn create_relation( }), } } + +pub fn create_relation_with_metadata( + source: String, + source_node_type: NodeType, + source_subtype: String, + to: String, + to_node_type: NodeType, + to_subtype: String, + rel_type: RelationType, + metadata: RelationMetadata, +) -> Relation { + let mut relation = + create_relation(source, source_node_type, source_subtype, to, to_node_type, to_subtype, rel_type); + relation.metadata = Some(metadata); + relation +} diff --git a/nucliadb_relations2/tests/test_reader.rs b/nucliadb_relations2/tests/test_reader.rs index 98064a416c..28fd90a1d9 100644 --- a/nucliadb_relations2/tests/test_reader.rs +++ b/nucliadb_relations2/tests/test_reader.rs @@ -24,8 +24,8 @@ use nucliadb_core::protos::entities_subgraph_request::DeletedEntities; use nucliadb_core::protos::relation::RelationType; use nucliadb_core::protos::relation_node::NodeType; use nucliadb_core::protos::{ - EntitiesSubgraphRequest, RelationNodeFilter, RelationPrefixSearchRequest, RelationSearchRequest, Resource, - ResourceId, + EntitiesSubgraphRequest, RelationMetadata, RelationNodeFilter, RelationPrefixSearchRequest, RelationSearchRequest, + Resource, ResourceId, }; use nucliadb_core::relations::*; use nucliadb_relations2::reader::RelationsReaderService; @@ -93,7 +93,7 @@ fn create_reader() -> NodeResult { "PEOPLE".to_string(), RelationType::Entity, ), - common::create_relation( + common::create_relation_with_metadata( "Anthony".to_string(), NodeType::Entity, "PEOPLE".to_string(), @@ -101,6 +101,15 @@ fn create_reader() -> NodeResult { NodeType::Entity, "PLACES".to_string(), RelationType::Entity, + RelationMetadata { + paragraph_id: Some("myresource/0/myresource/100-200".to_string()), + source_start: Some(0), + source_end: Some(10), + to_start: Some(11), + to_end: Some(20), + data_augmentation_task_id: Some("mytask".to_string()), + ..Default::default() + }, ), common::create_relation( "Anna".to_string(), @@ -212,6 +221,38 @@ fn test_search() -> NodeResult<()> { Ok(()) } +#[test] +fn test_search_metadata() -> NodeResult<()> { + let reader = create_reader()?; + + let result = reader.search(&RelationSearchRequest { + subgraph: Some(EntitiesSubgraphRequest { + depth: Some(1_i32), + entry_points: vec![common::create_relation_node( + "Anthony".to_string(), + NodeType::Entity, + "PEOPLE".to_string(), + )], + ..Default::default() + }), + ..Default::default() + })?; + + let subgraph = result.subgraph.unwrap(); + assert_eq!(subgraph.relations.len(), 1); + + let relation = &subgraph.relations[0]; + let metadata = relation.metadata.as_ref().unwrap(); + assert_eq!(metadata.paragraph_id, Some("myresource/0/myresource/100-200".to_string())); + assert_eq!(metadata.source_start, Some(0)); + assert_eq!(metadata.source_end, Some(10)); + assert_eq!(metadata.to_start, Some(11)); + assert_eq!(metadata.to_end, Some(20)); + assert_eq!(metadata.data_augmentation_task_id, Some("mytask".to_string())); + + Ok(()) +} + #[test] fn test_prefix_search() -> NodeResult<()> { let reader = create_reader()?; diff --git a/nucliadb_relations2/tests/test_writer.rs b/nucliadb_relations2/tests/test_writer.rs index 57934fa3fb..fa1e2b4d37 100644 --- a/nucliadb_relations2/tests/test_writer.rs +++ b/nucliadb_relations2/tests/test_writer.rs @@ -22,7 +22,7 @@ mod common; use nucliadb_core::prelude::*; use nucliadb_core::protos::relation::RelationType; use nucliadb_core::protos::relation_node::NodeType; -use nucliadb_core::protos::{Resource, ResourceId}; +use nucliadb_core::protos::{RelationMetadata, Resource, ResourceId}; use nucliadb_core::relations::*; use nucliadb_relations2::writer::RelationsWriterService; use tempfile::TempDir; @@ -69,7 +69,7 @@ fn test_index_docs() -> NodeResult<()> { shard_id: "shard_id".to_string(), }), relations: vec![ - common::create_relation( + common::create_relation_with_metadata( "cat".to_string(), NodeType::Entity, "ANIMALS".to_string(), @@ -77,6 +77,15 @@ fn test_index_docs() -> NodeResult<()> { NodeType::Entity, "ANIMALS".to_string(), RelationType::Entity, + RelationMetadata { + paragraph_id: Some("myresource/0/myresource/100-200".to_string()), + source_start: Some(0), + source_end: Some(10), + to_start: Some(11), + to_end: Some(20), + data_augmentation_task_id: Some("mytask".to_string()), + ..Default::default() + }, ), common::create_relation( "dolphin".to_string(),