From f401d9790ab4d9ef24549137a8a2967085661060 Mon Sep 17 00:00:00 2001
From: John Grimes <John.Grimes@csiro.au>
Date: Fri, 19 Apr 2024 09:24:09 +1000
Subject: [PATCH] Expand ExtractTest#simple and make assertions

---
 .../csiro/pathling/extract/ExtractTest.java   | 31 ++++++++++++++++---
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/fhirpath/src/test/java/au/csiro/pathling/extract/ExtractTest.java b/fhirpath/src/test/java/au/csiro/pathling/extract/ExtractTest.java
index 6dc4713183..53e81057e5 100644
--- a/fhirpath/src/test/java/au/csiro/pathling/extract/ExtractTest.java
+++ b/fhirpath/src/test/java/au/csiro/pathling/extract/ExtractTest.java
@@ -5,6 +5,8 @@
 import au.csiro.pathling.query.ExpressionWithLabel;
 import au.csiro.pathling.terminology.TerminologyServiceFactory;
 import au.csiro.pathling.test.SpringBootUnitTest;
+import au.csiro.pathling.test.assertions.DatasetAssert;
+import au.csiro.pathling.test.builders.DatasetBuilder;
 import au.csiro.pathling.test.datasource.ObjectDataSource;
 import ca.uhn.fhir.context.FhirContext;
 import java.util.Collections;
@@ -14,6 +16,7 @@
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.types.DataTypes;
 import org.hl7.fhir.r4.model.Enumerations.ResourceType;
 import org.hl7.fhir.r4.model.Patient;
 import org.junit.jupiter.api.Test;
@@ -39,9 +42,12 @@ class ExtractTest {
   void simple() {
     final Patient patient = new Patient();
     patient.setId("1");
-    patient.addName().setFamily("Smith").addGiven("John");
-    patient.addName().setFamily("Tyler").addGiven("Mary");
-    patient.addTelecom().setValue("555-555-5555");
+    patient.addName().setFamily("Kay").addGiven("Awee");
+    patient.addName().setFamily("Robert").addGiven("Zosia");
+    final Patient.ContactComponent contact1 = patient.addContact();
+    contact1.getName().setFamily("Ravindra").addGiven("Veremund");
+    final Patient.ContactComponent contact2 = patient.addContact();
+    contact2.getName().setFamily("Einion").addGiven("Kazuko");
     final ObjectDataSource dataSource = new ObjectDataSource(spark, encoders,
         List.of(patient));
     final ExtractQueryExecutor executor = new ExtractQueryExecutor(
@@ -52,11 +58,28 @@ void simple() {
         List.of(
             ExpressionWithLabel.of("id", "id"),
             ExpressionWithLabel.of("name.given", "given_name"),
-            ExpressionWithLabel.of("name.family", "family_name")
+            ExpressionWithLabel.of("name.family", "family_name"),
+            ExpressionWithLabel.of("contact.name.given", "contact_given_name"),
+            ExpressionWithLabel.of("contact.name.family", "contact_family_name")
         ), Collections.emptyList(), Optional.empty());
 
     final Dataset<Row> result = executor.buildQuery(request);
     result.show();
+
+    final Dataset<Row> expected = DatasetBuilder.of(spark)
+        .withColumn("id", DataTypes.StringType)
+        .withColumn("given_name", DataTypes.StringType)
+        .withColumn("family_name", DataTypes.StringType)
+        .withColumn("contact_given_name", DataTypes.StringType)
+        .withColumn("contact_family_name", DataTypes.StringType)
+        .withRow("1", "Awee", "Kay", "Veremund", "Ravindra")
+        .withRow("1", "Awee", "Kay", "Kazuko", "Einion")
+        .withRow("1", "Zosia", "Robert", "Veremund", "Ravindra")
+        .withRow("1", "Zosia", "Robert", "Kazuko", "Einion")
+        .build();
+
+    new DatasetAssert(result)
+        .hasRowsUnordered(expected);
   }
 
 }