Test refactor. address review comments

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/batch/processors/GoldenResourceAnnotatingProcessor.java

@@ -97,7 +97,8 @@ public class GoldenResourceAnnotatingProcessor implements ItemProcessor<List<IBa
 		if (patientReference.isPresent()) {
 			addGoldenResourceExtension(iBaseResource, patientReference.get());
 		} else {
-			ourLog.warn("Failed to find the patient reference information for resource {}", iBaseResource);
+			ourLog.error("Failed to find the patient reference information for resource {}. This is a bug, " +
+				"as all resources which can be exported via Group Bulk Export must reference a patient.", iBaseResource);
 		}
 	}
 

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/bulk/job/GroupBulkItemReader.java

@@ -36,6 +36,7 @@ import ca.uhn.fhir.model.primitive.IdDt;
 import ca.uhn.fhir.rest.api.server.storage.ResourcePersistentId;
 import ca.uhn.fhir.rest.param.ReferenceOrListParam;
 import ca.uhn.fhir.rest.param.ReferenceParam;
+import com.google.common.collect.Multimaps;
 import org.hl7.fhir.instance.model.api.IBaseResource;
 import org.hl7.fhir.instance.model.api.IIdType;
 import org.hl7.fhir.instance.model.api.IPrimitiveType;
@@ -147,15 +148,7 @@ public class GroupBulkItemReader extends BaseBulkItemReader implements ItemReade
 		//   patient/gold-2 -> [patient/3, patient/4]
 		//}
 		Map<Long, Set<Long>> goldenResourceToSourcePidMap = new HashMap<>();
-		for (IMdmLinkDao.MdmPidTuple goldenPidTargetPidTuple : thePidTuples) {
+		extract(thePidTuples, goldenResourceToSourcePidMap);
-			Long goldenPid = goldenPidTargetPidTuple.getGoldenPid();
-			Long sourcePid = goldenPidTargetPidTuple.getSourcePid();
-
-			if(!goldenResourceToSourcePidMap.containsKey(goldenPid)) {
-				goldenResourceToSourcePidMap.put(goldenPid, new HashSet<>());
-			}
-			goldenResourceToSourcePidMap.get(goldenPid).add(sourcePid);
-		}
 
 		//Next, lets convert it to an inverted index for fast lookup
 		// {
@@ -217,15 +210,7 @@ public class GroupBulkItemReader extends BaseBulkItemReader implements ItemReade
 			Map<Long, Optional<String>> pidToForcedIdMap = myIdHelperService.translatePidsToForcedIds(uniquePids);
 
 			Map<Long, Set<Long>> goldenResourceToSourcePidMap = new HashMap<>();
-			for (IMdmLinkDao.MdmPidTuple goldenPidTargetPidTuple : goldenPidTargetPidTuples) {
+			extract(goldenPidTargetPidTuples, goldenResourceToSourcePidMap);
-				Long goldenPid = goldenPidTargetPidTuple.getGoldenPid();
-				Long sourcePid = goldenPidTargetPidTuple.getSourcePid();
-
-				if(!goldenResourceToSourcePidMap.containsKey(goldenPid)) {
-					goldenResourceToSourcePidMap.put(goldenPid, new HashSet<>());
-				}
-				goldenResourceToSourcePidMap.get(goldenPid).add(sourcePid);
-			}
 			populateMdmResourceCache(goldenPidTargetPidTuples);
 
 			//If the result of the translation is an empty optional, it means there is no forced id, and we can use the PID as the resource ID.
@@ -243,6 +228,14 @@ public class GroupBulkItemReader extends BaseBulkItemReader implements ItemReade
 		return expandedIds;
 	}
 
+	private void extract(List<IMdmLinkDao.MdmPidTuple> theGoldenPidTargetPidTuples, Map<Long, Set<Long>> theGoldenResourceToSourcePidMap) {
+		for (IMdmLinkDao.MdmPidTuple goldenPidTargetPidTuple : theGoldenPidTargetPidTuples) {
+			Long goldenPid = goldenPidTargetPidTuple.getGoldenPid();
+			Long sourcePid = goldenPidTargetPidTuple.getSourcePid();
+			theGoldenResourceToSourcePidMap.computeIfAbsent(goldenPid, key -> new HashSet<>()).add(sourcePid);
+		}
+	}
+
 	private void queryResourceTypeWithReferencesToPatients(Set<ResourcePersistentId> myReadPids, List<String> idChunk) {
 		//Build SP map
 		//First, inject the _typeFilters and _since from the export job

hapi-fhir-jpaserver-base/src/test/java/ca/uhn/fhir/jpa/bulk/BulkDataExportSvcImplR4Test.java

@@ -21,6 +21,7 @@ import ca.uhn.fhir.jpa.entity.BulkExportJobEntity;
 import ca.uhn.fhir.jpa.entity.MdmLink;
 import ca.uhn.fhir.mdm.api.MdmLinkSourceEnum;
 import ca.uhn.fhir.mdm.api.MdmMatchResultEnum;
+import ca.uhn.fhir.parser.IParser;
 import ca.uhn.fhir.rest.api.Constants;
 import ca.uhn.fhir.rest.server.exceptions.InvalidRequestException;
 import ca.uhn.fhir.util.HapiExtensions;
@@ -29,11 +30,13 @@ import com.google.common.base.Charsets;
 import com.google.common.collect.Sets;
 import org.apache.commons.lang3.time.DateUtils;
 import org.hamcrest.Matchers;
+import org.hl7.fhir.instance.model.api.IBaseResource;
 import org.hl7.fhir.instance.model.api.IIdType;
 import org.hl7.fhir.r4.model.Binary;
 import org.hl7.fhir.r4.model.CareTeam;
 import org.hl7.fhir.r4.model.CodeableConcept;
 import org.hl7.fhir.r4.model.Enumerations;
+import org.hl7.fhir.r4.model.Extension;
 import org.hl7.fhir.r4.model.Group;
 import org.hl7.fhir.r4.model.Immunization;
 import org.hl7.fhir.r4.model.InstantType;
@@ -643,7 +646,7 @@ public class BulkDataExportSvcImplR4Test extends BaseJpaR4Test {
 		// Create a bulk job
 		BulkDataExportOptions bulkDataExportOptions = new BulkDataExportOptions();
 		bulkDataExportOptions.setOutputFormat(null);
-		bulkDataExportOptions.setResourceTypes(Sets.newHashSet("Immunization", "Observation", "Patient"));
+		bulkDataExportOptions.setResourceTypes(Sets.newHashSet("Immunization", "Patient"));
 		bulkDataExportOptions.setSince(null);
 		bulkDataExportOptions.setFilters(null);
 		bulkDataExportOptions.setGroupId(myPatientGroupId);
@@ -660,19 +663,36 @@ public class BulkDataExportSvcImplR4Test extends BaseJpaR4Test {
 		assertThat(jobInfo.getFiles().size(), equalTo(3));
 		assertThat(jobInfo.getFiles().get(0).getResourceType(), is(equalTo("Immunization")));
 
-		// Iterate over the files
+		//Ensure that all immunizations refer to the golden resource via extension
-
-		String nextContents = getBinaryContents(jobInfo, 0);
 		assertThat(jobInfo.getFiles().get(0).getResourceType(), is(equalTo("Immunization")));
-		assertThat(nextContents, is(containsString(HapiExtensions.ASSOCIATED_GOLDEN_RESOURCE_EXTENSION_URL)));
+		List<Immunization> immunizations = readBulkExportContentsIntoResources(getBinaryContents(jobInfo, 0), Immunization.class);
+		immunizations
+			.stream().filter(immu -> !immu.getIdElement().getIdPart().equals("PAT999"))//Skip the golden resource
+			.forEach(immunization -> {
+				Extension extensionByUrl = immunization.getExtensionByUrl(HapiExtensions.ASSOCIATED_GOLDEN_RESOURCE_EXTENSION_URL);
+				String reference = ((Reference) extensionByUrl.getValue()).getReference();
+				assertThat(reference, is(equalTo("Patient/PAT999")));
+			});
 
-		nextContents = getBinaryContents(jobInfo, 1);
+		//Ensure all patients are linked to their golden resource.
-		assertThat(jobInfo.getFiles().get(1).getResourceType(), is(equalTo("Observation")));
+		assertThat(jobInfo.getFiles().get(1).getResourceType(), is(equalTo("Patient")));
-		assertThat(nextContents, is(containsString(HapiExtensions.ASSOCIATED_GOLDEN_RESOURCE_EXTENSION_URL)));
+		List<Patient> patients = readBulkExportContentsIntoResources(getBinaryContents(jobInfo, 2), Patient.class);
+		patients.stream()
+			.filter(patient -> patient.getIdElement().getIdPart().equals("PAT999"))
+			.forEach(patient -> {
+				Extension extensionByUrl = patient.getExtensionByUrl(HapiExtensions.ASSOCIATED_GOLDEN_RESOURCE_EXTENSION_URL);
+				String reference = ((Reference) extensionByUrl.getValue()).getReference();
+				assertThat(reference, is(equalTo("Patient/PAT999")));
+			});
 
-		nextContents = getBinaryContents(jobInfo, 2);
+	}
-		assertThat(jobInfo.getFiles().get(2).getResourceType(), is(equalTo("Patient")));
+
-		assertThat(nextContents, is(containsString(HapiExtensions.ASSOCIATED_GOLDEN_RESOURCE_EXTENSION_URL)));
+	private <T extends IBaseResource> List<T> readBulkExportContentsIntoResources(String theContents, Class<T> theClass) {
+		IParser iParser = myFhirCtx.newJsonParser();
+		return Arrays.stream(theContents.split("\n"))
+			.map(iParser::parseResource)
+			.map(theClass::cast)
+			.collect(Collectors.toList());
 	}
 
 	@Test