Add unit tests

nuclia · Jan 9, 2025 · f08312e · f08312e
1 parent 3f9e22c
commit f08312e
Show file tree

Hide file tree

Showing 4 changed files with 87 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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,14 +93,23 @@ fn create_reader() -> NodeResult<RelationsReaderService> {
                 "PEOPLE".to_string(),
                 RelationType::Entity,
             ),
-            common::create_relation(
+            common::create_relation_with_metadata(
                 "Anthony".to_string(),
                 NodeType::Entity,
                 "PEOPLE".to_string(),
                 "Netherlands".to_string(),
                 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
@@ -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,14 +69,23 @@ 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(),
                 "cat".to_string(),
                 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(),