Merge branch 'extending-tests' into 'master'

Extend Test Cases Phase 1

See merge request JDOsborne1/db_to_d2!2

Author: JDOsborne1

cmd/configuration.go

@@ -4,6 +4,7 @@ import (
 	"core"
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"virtual"
 )
@@ -12,30 +13,58 @@ import (
 // These are currently set by a file specified by the VIRTUAL_LINKS_PATH environment variable.
 // The file should be a json array of virtual links. See the virtual package for more information.
 func get_virtual_links() []virtual.VirtualLink {
-	links := []virtual.VirtualLink{}
-	links_json, err := os.ReadFile(os.Getenv("VIRTUAL_LINKS_PATH"))
+	links_reader, err := os.Open(os.Getenv("VIRTUAL_LINKS_PATH"))
 	if err != nil {
 		//TODO: Log error, or bubble up instead of printing to console
+		fmt.Println("Failed to open virtual links file")
+	}
+	links, err := read_virtual_links(links_reader)
+	if err != nil {
 		fmt.Println("Failed to read virtual links file")
 	}
-	json.Unmarshal(links_json, &links)
+
 	return links
 }
 
+func read_virtual_links(_input io.Reader) ([]virtual.VirtualLink, error) {
+	links := []virtual.VirtualLink{}
+	links_json, err := io.ReadAll(_input)
+	if err != nil {
+		return links, err
+	}
+	err = json.Unmarshal(links_json, &links)
+	return links, err
+}
+
 // get_table_groups returns the table groups for the program.
 // These are currently set by a file specified by the TABLE_GROUPS_PATH environment variable.
 // The file should be a json array of table groups. See the core package for more information.
 func get_table_groups() []core.TableGroup {
 	table_groups := []core.TableGroup{}
-	table_groups_json, err := os.ReadFile(os.Getenv("TABLE_GROUPS_PATH"))
+	table_groups_reader, err := os.Open(os.Getenv("TABLE_GROUPS_PATH"))
 	if err != nil {
 		//TODO: Log error, or bubble up instead of printing to console
 		fmt.Println("Failed to read table groups file")
 	}
-	json.Unmarshal(table_groups_json, &table_groups)
+
+	table_groups, err = read_table_groups(table_groups_reader)
+
+	if err != nil {
+		fmt.Println("Failed to read table groups file")
+	}
 	return table_groups
 }
 
+func read_table_groups(_input io.Reader) ([]core.TableGroup, error) {
+	table_groups := []core.TableGroup{}
+	table_groups_json, err := io.ReadAll(_input)
+	if err != nil {
+		return table_groups, err
+	}
+	err = json.Unmarshal(table_groups_json, &table_groups)
+	return table_groups, err
+}
+
 // get_designated_user returns the designated user for the program. Set by the DESIGNATED_USER environment variable.
 // This is used to restrict the schema to the tables that the designated user has access to. See the mysql package for more information.
 func get_designated_user() string {

pkg/core/d2_generator_test.go

@@ -0,0 +1,74 @@
+package core
+
+// The d2 tests can be done by using strings.TrimSpace() to remove the whitespace from the output of the d2 generator.
+// This should remove the influence of formatting differences between the two strings. This is a simple way to test that that two are equal.
+
+import (
+	"strings"
+	"testing"
+
+)
+
+func simple_formatter(_input string) string {
+	_input = strings.Replace(_input, "\t", "", -1)
+	_input = strings.Replace(_input, "\n\n", "\n", -1)
+	_input = strings.TrimSpace(_input)
+	return _input 
+}
+
+func Test_d2_generator(t *testing.T) {
+	schema := Schema{
+		Tables: []Table{
+			{
+				Name: "table1",
+				Columns: []Column{
+					{
+						Name: "col1",
+						Type: "int",
+					},
+					{
+						Name: "col2",
+						Type: "varchar",
+					},
+				},
+			},
+			{
+				Name: "table2",
+				Columns: []Column{
+					{
+						Name: "col1",
+						Type: "int",
+					},
+					{
+						Name: "col2",
+						Type: "varchar",
+					},
+				},
+			},
+		},
+	}
+
+expected := `
+table1: {
+  		shape: sql_table
+  col1: int
+  		col2: varchar
+}
+
+table2: {
+  shape: sql_table
+  col1: int
+
+  col2: varchar
+}
+`
+	expected = simple_formatter(expected)
+
+	actual := simple_formatter(Schema_to_d2(schema, []TableGroup{}))
+
+	if actual != expected {
+		t.Errorf("Expected \n%s\n, got \n%s", expected, actual)
+	}
+
+}
+

pkg/core/go.sum

@@ -5,29 +5,175 @@ import (
 	"testing"
 )
 
-// TestRestrictionsIdentity is a basic test that checks that the identity function returns the same value as the input.
-func TestRestrictionsIdentity(t *testing.T) {
-	//Test data
-	input := Schema{
-		Tables: []Table{
-			{
-				Name: "test",
-				Columns: []Column{
-					{
-						Name: "test",
+var example_schema = Schema{
+	Tables: []Table{
+		{
+			Name: "users",
+			Columns: []Column{
+				{
+					Name:     "id",
+					Type:     "int",
+					Nullable: false,
+					Key:      "PRIMARY KEY",
+					Extra:    "AUTO_INCREMENT",
+				},
+				{
+					Name:     "username",
+					Type:     "varchar(255)",
+					Nullable: false,
+					Key:      "UNIQUE KEY",
+					Extra:    "",
+				},
+				{
+					Name:     "email",
+					Type:     "varchar(255)",
+					Nullable: false,
+					Key:      "UNIQUE KEY",
+					Extra:    "",
+				},
+			},
+		},
+		{
+			Name: "posts",
+			Columns: []Column{
+				{
+					Name:     "id",
+					Type:     "int",
+					Nullable: false,
+					Key:      "PRIMARY KEY",
+					Extra:    "AUTO_INCREMENT",
+				},
+				{
+					Name:     "title",
+					Type:     "varchar(255)",
+					Nullable: false,
+					Key:      "",
+					Extra:    "",
+				},
+				{
+					Name:     "content",
+					Type:     "text",
+					Nullable: true,
+					Key:      "",
+					Extra:    "",
+				},
+				{
+					Name:     "user_id",
+					Type:     "int",
+					Nullable: false,
+					Key:      "FOREIGN KEY",
+					Extra:    "",
+					Reference: &Reference{
+						Table:    "users",
+						Column:   "id",
+						OnDelete: "CASCADE",
+						OnUpdate: "CASCADE",
 					},
 				},
 			},
 		},
+		{
+			Name: "meta",
+			Columns: []Column{
+				{
+					Name:     "meta_key",
+					Type:     "varchar(255)",
+					Nullable: false,
+					Key:      "",
+					Extra:    "",
+				},
+				{
+					Name:     "meta_value",
+					Type:     "varchar(255)",
+					Nullable: false,
+					Key:      "",
+					Extra:    "",
+				},
+			},
+		},
+	},
+}
+
+func get_table_names(_input Schema) []string {
+	var output []string
+	for _, table := range _input.Tables {
+		output = append(output, table.Name)
+	}
+	return output
+}
+
+func get_column_names(_input Table) []string {
+	var output []string
+	for _, column := range _input.Columns {
+		output = append(output, column.Name)
 	}
-	expected := input
+	return output
+}
 
+// TestRestrictionsIdentity is a basic test that checks that the identity function returns the same value as the input.
+func TestRestrictionsIdentity(t *testing.T) {
 	//Execute test
-	actual := Restrict(input, Standard)
+	actual := Restrict(example_schema, Standard)
 
 	//Compare actual to expected
-	if !reflect.DeepEqual(actual, expected) {
+	if !reflect.DeepEqual(actual, example_schema) {
 		t.Log("Identity function failed to return the same value as the input.")
 		t.Fail()
 	}
 }
+
+func TestRestrictionsMinimalist(t *testing.T) {
+	actual := Restrict(example_schema, Minimalist)
+	names := get_table_names(actual)
+	if !(equal_set(names, []string{"users", "posts"})) {
+		t.Log("Minimalist cleared inappropriate tables.")
+		t.Fail()
+	}
+
+	if !(equal_set(get_column_names(actual.Tables[0]), []string{"id", "username", "email"})) {
+		t.Log("Minimalist cleared columns in a table where all the columns are keys.")
+		t.Fail()
+	}
+
+	if !(equal_set(get_column_names(actual.Tables[1]), []string{"id", "user_id"})) {
+		t.Log("Minimalist didn't clear the right columns in a table which has partial keys.")
+		t.Fail()
+	}
+
+}
+
+// Simple example of a restriction function. This is intended to simulate the scenario where a user has been 
+// restricted to not be able to see the PII (Personally Identifiable Information) of other users. But they can 
+// still access the user_id fields for analysis purposes.
+func example_permission_restrictor_analyst(_table Table, _column Column) bool {
+	return _column.Name == "username" || _column.Name == "email"
+}
+
+// Another example of a restriction function. This is intended to simulate the scenario where a user has been
+// restricted to not be able to see the content generated by users, but can still edit the metadata associated
+// with the a user. This could be an example permission profile for a client account for the user management service.
+func example_permission_restrictor_user_profile(_table Table, _column Column) bool {
+	return _table.Name == "posts"
+}
+
+func TestRestrictionsCustom(t *testing.T) {
+	analyst := Restrict(example_schema, example_permission_restrictor_analyst)
+	if !(equal_set(get_table_names(analyst), []string{"users", "posts", "meta"})) {
+		t.Log("Custom cleared inappropriate tables.")
+		t.Fail()
+	}
+
+	if !(equal_set(get_column_names(analyst.Tables[0]), []string{"id"})) {
+		t.Log("Custom failed to clear the PII columns in the users table.")
+		t.Fail()
+	}
+
+	profile_service := Restrict(example_schema, example_permission_restrictor_user_profile)
+
+	if !(equal_set(get_table_names(profile_service), []string{"users","meta"})) {
+		t.Log("Custom cleared inappropriate tables.")
+		t.Fail()
+	}
+
+	
+}

pkg/core/restrictions_test.go

@@ -9,6 +9,18 @@ func in_set(_element string, _set []string) bool {
 	return false
 }
 
+func equal_set(_set1 []string, _set2 []string) bool {
+	if len(_set1) != len(_set2) {
+		return false
+	}
+	for _, element := range _set1 {
+		if !in_set(element, _set2) {
+			return false
+		}
+	}
+	return true
+}
+
 // Wraps a table name in a group tag if it is in a group. Otherwise, returns the table name.
 // This is used to ensure that links in the d2 graph are drawn correctly when table groups are used.
 func wrap_name_in_group(_table_name string, _grouping []TableGroup) string {