diff --git a/tests/data_generator/datagenerator.go b/tests/data_generator/datagenerator.go
index b3dcc9f..2449918 100644
--- a/tests/data_generator/datagenerator.go
+++ b/tests/data_generator/datagenerator.go
@@ -325,6 +325,7 @@ func main() {
 	environment := flag.String("e", "development", "Environment/prefix of config .yaml file name")
 	testData := flag.String("d", "models_tests_data_config", "Prefix of test data config .yaml file name")
 	sourceId := flag.Int("s", 1, "Source id for Omop DB")
+	tests.SetTestSourceId(*sourceId)
 	flag.Parse()
 	Init(*environment, *sourceId)
 	log.Printf("\n\n=============== GENERATING TEST DATA BASED ON CONFIG ============================\n\n")
diff --git a/tests/testutils.go b/tests/testutils.go
index 467d752..c100250 100644
--- a/tests/testutils.go
+++ b/tests/testutils.go
@@ -17,8 +17,15 @@ import (
 	"gorm.io/gorm"
 )
 
+// Global variable
+var testSourceId int = 1
+
 func GetTestSourceId() int {
-	return 1 // TODO - ideally this should also be used when populating "source" tables in test Atlas DB in the first place...
+	return testSourceId
+}
+
+func SetTestSourceId(id int) {
+	testSourceId = id
 }
 
 func GetTestDummyContinuousConceptId() int64 {