tools: fix addIndividual function, #TASK-7170

opencb · Nov 8, 2024 · bc0f8bc · bc0f8bc
1 parent 00eb026
commit bc0f8bc
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 8 deletions.
diff --git a/...tools/src/main/java/org/opencb/biodata/tools/variant/metadata/VariantMetadataManager.java b/...tools/src/main/java/org/opencb/biodata/tools/variant/metadata/VariantMetadataManager.java
@@ -35,7 +35,6 @@
 import org.opencb.biodata.models.variant.metadata.VariantFileMetadata;
 import org.opencb.biodata.models.variant.metadata.VariantMetadata;
 import org.opencb.biodata.models.variant.metadata.VariantStudyMetadata;
-import org.opencb.biodata.tools.variant.converters.avro.VCFHeaderToVariantFileHeaderConverter;
 import org.opencb.biodata.tools.variant.converters.avro.VCFHeaderToVariantFileMetadataConverter;
 import org.opencb.commons.datastore.core.Query;
 import org.opencb.commons.utils.FileUtils;
@@ -59,6 +58,7 @@
 public class VariantMetadataManager {
 
     private VariantMetadata variantMetadata;
+    private Set<String> individualSet = new HashSet<>();
 
     private ObjectMapper mapper;
     private Logger logger;
@@ -109,6 +109,7 @@ public void load(Path path) throws IOException {
         for (VariantStudyMetadata variantStudyMetadata: variantMetadata.getStudies()) {
             if (variantStudyMetadata.getIndividuals() != null) {
                 for (Individual individual : variantStudyMetadata.getIndividuals()) {
+                    individualSet.add(individual.getId());
                     for (Sample sample : individual.getSamples()) {
                         sample.getAnnotations().put(INDIVIDUAL_ID, individual.getId());
                         sample.getAnnotations().put(INDIVIDUAL_FAMILY, individual.getFamily());
@@ -272,6 +273,7 @@ public void addFile(VariantFileMetadata fileMetadata, String studyId) {
             individual.setSamples(samples);
 
             variantStudyMetadata.getIndividuals().add(individual);
+            individualSet.add(individual.getId());
         }
 
 
@@ -383,14 +385,12 @@ public void addIndividual(Individual individual, String studyId) {
         if (variantStudyMetadata.getIndividuals() == null) {
             variantStudyMetadata.setIndividuals(new ArrayList<>());
         }
-        for (Individual indi: variantStudyMetadata.getIndividuals()) {
-            if (indi.getId() != null && indi.getId().equals(individual.getId())) {
-                logger.error("Individual with id '{}' already exists in study '{}'", individual.getId(),
-                        studyId);
-                return;
-            }
+        if (individualSet.contains(individual.getId())) {
+            logger.error("Individual with id '{}' already exists in study '{}'", individual.getId(), studyId);
+            return;
         }
         variantStudyMetadata.getIndividuals().add(individual);
+        individualSet.add(individual.getId());
     }
 
     /**
@@ -430,6 +430,7 @@ public void removeIndividual(String individualId, String studyId) {
             for (int i = 0; i < variantStudyMetadata.getIndividuals().size(); i++) {
                 if (individualId.equals(variantStudyMetadata.getIndividuals().get(i).getId())) {
                     variantStudyMetadata.getIndividuals().remove(i);
+                    individualSet.remove(individualId);
                     return;
                 }
             }

diff --git a/...s/src/test/java/org/opencb/biodata/tools/variant/metadata/VariantMetadataManagerTest.java b/...s/src/test/java/org/opencb/biodata/tools/variant/metadata/VariantMetadataManagerTest.java
@@ -1,13 +1,13 @@
 package org.opencb.biodata.tools.variant.metadata;
 
 import org.apache.commons.collections4.map.HashedMap;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.opencb.biodata.models.metadata.Individual;
 import org.opencb.biodata.models.metadata.Sample;
 import org.opencb.biodata.models.variant.metadata.VariantMetadata;
 import org.opencb.biodata.models.variant.metadata.VariantStudyMetadata;
-import org.opencb.biodata.tools.variant.metadata.VariantMetadataManager;
 import org.opencb.commons.datastore.core.Query;
 
 import java.io.IOException;
@@ -184,4 +184,38 @@ public void load() {
         }
     }
 
+    @Test
+    public void testAddIndividual() {
+        try {
+            Path path = Paths.get("/tmp/ds.meta.json");
+            if (path.toFile().exists()) {
+                path.toFile().delete();
+            }
+            manager.save(path, true);
+
+            manager = new VariantMetadataManager();
+            manager.load(path);
+
+            String studyId = variantMetadata.getStudies().get(0).getId();
+            Assert.assertEquals(3, manager.getVariantStudyMetadata(studyId).getIndividuals().size());
+
+            manager.addIndividual("Person_0", "Sample_0_0", studyId);
+
+            String newIndividualId = "Person_100";
+            manager.addIndividual(newIndividualId, "Sample_100_0", studyId);
+            Assert.assertEquals(4, manager.getVariantStudyMetadata(studyId).getIndividuals().size());
+
+            manager.addIndividual(newIndividualId, "Sample_100_0", studyId);
+            Assert.assertEquals(4, manager.getVariantStudyMetadata(studyId).getIndividuals().size());
+
+            manager.removeIndividual(newIndividualId, studyId);
+            Assert.assertEquals(3, manager.getVariantStudyMetadata(studyId).getIndividuals().size());
+
+            manager.addIndividual(newIndividualId, "Sample_100_0", studyId);
+            Assert.assertEquals(4, manager.getVariantStudyMetadata(studyId).getIndividuals().size());
+        } catch (IOException e) {
+            e.printStackTrace();
+            fail();
+        }
+    }
 }