From ac207fbe3a92b3f0a7d9f5212ec8aa494fa1c1e2 Mon Sep 17 00:00:00 2001
From: schuemie <schuemie@ohdsi.org>
Date: Tue, 19 Jul 2016 11:04:25 +0200
Subject: [PATCH] Added ability to look up values in other tables in the target
 schema when defining tests.

---
 .../ETLTestFrameWorkGenerator.java            | 47 +++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/src/org/ohdsi/rabbitInAHat/ETLTestFrameWorkGenerator.java b/src/org/ohdsi/rabbitInAHat/ETLTestFrameWorkGenerator.java
index 8c15ade4..33081a2e 100644
--- a/src/org/ohdsi/rabbitInAHat/ETLTestFrameWorkGenerator.java
+++ b/src/org/ohdsi/rabbitInAHat/ETLTestFrameWorkGenerator.java
@@ -74,6 +74,7 @@ private static List<String> generateRScript(ETL etl) {
 		createExpectFunctions(r, DEFAULT, etl.getTargetDatabase());
 		createExpectFunctions(r, NEGATE, etl.getTargetDatabase());
 		createExpectFunctions(r, COUNT, etl.getTargetDatabase());
+		createLookupFunctions(r, etl.getTargetDatabase());
 		return r;
 	}
 
@@ -107,6 +108,8 @@ private static void createExpectFunctions(List<String> r, int type, Database dat
 				testDefs.add("    }");
 				testDefs.add("    if (is.null(" + rFieldName + ")) {");
 				testDefs.add("      statement <- paste0(statement, \" " + sqlFieldName + " IS NULL\")");
+				testDefs.add("    } else if (is(" + rFieldName + ", \"subQuery\")){");
+				testDefs.add("      statement <- paste0(statement, \" " + sqlFieldName + " = (\", as.character(" + rFieldName + "), \")\")");
 				testDefs.add("    } else {");
 				testDefs.add("      statement <- paste0(statement, \" " + sqlFieldName + " = '\", " + rFieldName + ",\"'\")");
 				testDefs.add("    }");
@@ -150,6 +153,50 @@ else if (type == NEGATE)
 			r.add("");
 		}
 	}
+	
+	private static void createLookupFunctions(List<String> r, Database database) {
+		for (Table table : database.getTables()) {
+			StringBuilder line = new StringBuilder();
+			String rTableName = convertToRName(table.getName());
+			String sqlTableName = convertToSqlName(table.getName());
+			List<String> argDefs = new ArrayList<String>();
+			List<String> testDefs = new ArrayList<String>();
+			for (Field field : table.getFields()) {
+				String rFieldName = convertToRName(field.getName());
+				String sqlFieldName = convertToSqlName(field.getName());
+				argDefs.add(rFieldName);
+				testDefs.add("  if (!missing(" + rFieldName + ")) {");
+				testDefs.add("    if (first) {");
+				testDefs.add("      first <- FALSE");
+				testDefs.add("    } else {");
+				testDefs.add("      statement <- paste0(statement, \" AND\")");
+				testDefs.add("    }");
+				testDefs.add("    if (is.null(" + rFieldName + ")) {");
+				testDefs.add("      statement <- paste0(statement, \" " + sqlFieldName + " IS NULL\")");
+				testDefs.add("    } else {");
+				testDefs.add("      statement <- paste0(statement, \" " + sqlFieldName + " = '\", " + rFieldName + ",\"'\")");
+				testDefs.add("    }");
+				testDefs.add("  }");
+				testDefs.add("");
+			}
+			line.append("lookup_" + rTableName + " <- function(fetchField, ");
+			line.append(StringUtilities.join(argDefs, ", "));
+			line.append(") {");
+			r.add(line.toString());
+
+			line = new StringBuilder();
+			line.append("  statement <- paste0(\"SELECT \", fetchField , \" FROM ");
+			line.append(sqlTableName);
+			line.append(" WHERE\")");
+			r.add(line.toString());
+			r.add("  first <- TRUE");
+			r.addAll(testDefs);
+			r.add("  class(statement) <- \"subQuery\"");
+			r.add("  return(statement)");
+			r.add("}");
+			r.add("");
+		}
+	}
 
 	private static String convertToSqlName(String name) {
 		if (name.contains(" ") || name.contains(".") || keywordSet.contains(name.toUpperCase()))