3 files changed, 38 insertions, 5 deletions
diff --git a/pkg/datamaps/db.go b/pkg/datamaps/db.go
index e3b5e7a..4c3605c 100644
--- a/pkg/datamaps/db.go
+++ b/pkg/datamaps/db.go
@@ -15,6 +15,9 @@ import (
 // setupDB creates the intitial database
 func setupDB(path string) (*sql.DB, error) {
 	stmtBase := `DROP TABLE IF EXISTS datamap;
+				 DROP TABLE IF EXISTS datamap_line;
+				 DROP TABLE IF EXISTS return;
+				 DROP TABLE IF EXISTS return_data;
 				  CREATE TABLE datamap(id INTEGER PRIMARY KEY, name TEXT, date_created TEXT);
 				  DROP TABLE IF EXISTS datamap_line;
 
@@ -28,6 +31,19 @@ func setupDB(path string) (*sql.DB, error) {
 					REFERENCES datamap(id) 
 					ON DELETE CASCADE      
 				  );                        
+				 CREATE TABLE return(
+					 id INTEGER PRIMARY KEY,
+					 name TEXT,
+					 data_created TEXT
+					);
+				 CREATE TABLE return_data(
+					 id INTEGER PRIMARY KEY,
+					 dml_id INTEGER,
+					 value TEXT,
+					 FOREIGN KEY (dml_id)
+					 REFERENCES datamap_line(id) 
+					 ON DELETE CASCADE
+				 );
 				 `
 	os.Create(path)
 	db, err := sql.Open("sqlite3", path)
diff --git a/pkg/datamaps/db_test.go b/pkg/datamaps/db_test.go
index 2055d8a..3ca9054 100644
--- a/pkg/datamaps/db_test.go
+++ b/pkg/datamaps/db_test.go
@@ -1,12 +1,16 @@
 package datamaps
 
 import (
+	"os"
 	"testing"
 )
 
 func TestOpenSQLiteFile(t *testing.T) {
 	db, err := setupDB("./testdata/test.db")
-	defer db.Close()
+	defer func() {
+		db.Close()
+		os.Remove("./testdata/test.db")
+	}()
 
 	if err != nil {
 		t.Fatal("Expected to be able to set up the database.")
@@ -38,7 +42,10 @@ func TestOpenSQLiteFile(t *testing.T) {
 
 func TestDatamapGoesIntoDB(t *testing.T) {
 	db, err := setupDB("./testdata/test.db")
-	defer db.Close()
+	defer func() {
+		db.Close()
+		os.Remove("./testdata/test.db")
+	}()
 
 	if err != nil {
 		t.Fatal("Expected to be able to set up the database.")
diff --git a/pkg/datamaps/reader_test.go b/pkg/datamaps/reader_test.go
index c221b26..53aff4a 100644
--- a/pkg/datamaps/reader_test.go
+++ b/pkg/datamaps/reader_test.go
@@ -1,6 +1,7 @@
 package datamaps
 
 import (
+	"os"
 	"testing"
 )
 
@@ -73,7 +74,10 @@ func TestReadXLSX(t *testing.T) {
 // func TestExtractWithDBDatamap(t *testing.T) {
 // 	// setup - we need the datamap in the test database
 // 	db, err := setupDB("./testdata/test.db")
-// 	defer db.Close()
+// defer func() {
+// 	db.Close()
+// 	os.Remove("./testdata/test.db")
+// }()
 
 // 	if err != nil {
 // 		t.Fatal("Expected to be able to set up the database.")
@@ -95,7 +99,10 @@ func TestReadXLSX(t *testing.T) {
 func TestDMLSliceFromDatabase(t *testing.T) {
 	// setup - we need the datamap in the test database
 	db, err := setupDB("./testdata/test.db")
-	defer db.Close()
+	defer func() {
+		db.Close()
+		os.Remove("./testdata/test.db")
+	}()
 
 	if err != nil {
 		t.Fatal("Expected to be able to set up the database.")
@@ -142,7 +149,10 @@ func TestDMLSliceFromDatabase(t *testing.T) {
 func TestExtractUsingDBDM(t *testing.T) {
 	// setup - we need the datamap in the test database
 	db, err := setupDB("./testdata/test.db")
-	defer db.Close()
+	defer func() {
+		db.Close()
+		os.Remove("./testdata/test.db")
+	}()
 
 	if err != nil {
 		t.Fatal("Expected to be able to set up the database.")