diff --git a/data/TEST/importedData.csv b/data/TEST/importedData.csv
index c2c1827..2e08817 100644
--- a/data/TEST/importedData.csv
+++ b/data/TEST/importedData.csv
@@ -118,7 +118,7 @@ DELAY,PAYED,NET-WORTH,INFLUENCE,SKARBOWKA,MEMBER,HAT,SIZE,PRIORITY
 8,TRUE,44,95,FALSE,TRUE,FALSE,CompanySize.BIG,MEDIUM
 0,TRUE,40,3,FALSE,FALSE,FALSE,CompanySize.SMALL,LOW
 2,TRUE,49,39,FALSE,TRUE,TRUE,CompanySize.NO,LOW
-10,TRUE,4,94,FALSE,TRUE,TRUE,CompanySize.BIG,HIGH
+10,TRUE,4,94,FALSE,TRUE,TRUE,CompanySize.BIG,MEDIUM
 0,TRUE,90,86,FALSE,TRUE,FALSE,CompanySize.SMALL,MEDIUM
 14,TRUE,61,74,FALSE,FALSE,FALSE,CompanySize.GIGANTISHE,MEDIUM
 3,TRUE,86,27,FALSE,FALSE,TRUE,CompanySize.NORMAL,MEDIUM
@@ -127,7 +127,7 @@ DELAY,PAYED,NET-WORTH,INFLUENCE,SKARBOWKA,MEMBER,HAT,SIZE,PRIORITY
 13,TRUE,96,89,FALSE,FALSE,FALSE,CompanySize.NO,MEDIUM
 7,TRUE,83,73,FALSE,TRUE,TRUE,CompanySize.HUGE,HIGH
 7,TRUE,21,30,FALSE,FALSE,TRUE,CompanySize.NO,LOW
-10,TRUE,26,41,FALSE,TRUE,TRUE,CompanySize.NO,MEDIUM
+10,TRUE,26,41,FALSE,TRUE,TRUE,CompanySize.NO,LOW
 13,FALSE,73,16,FALSE,TRUE,FALSE,CompanySize.NO,LOW
 12,TRUE,30,54,TRUE,TRUE,TRUE,CompanySize.GIGANTISHE,HIGH
 12,FALSE,9,52,FALSE,FALSE,FALSE,CompanySize.NORMAL,LOW
@@ -180,7 +180,7 @@ DELAY,PAYED,NET-WORTH,INFLUENCE,SKARBOWKA,MEMBER,HAT,SIZE,PRIORITY
 5,FALSE,24,41,FALSE,FALSE,FALSE,CompanySize.BIG,LOW
 0,TRUE,0,25,FALSE,TRUE,FALSE,CompanySize.BIG,MEDIUM
 13,TRUE,41,91,TRUE,FALSE,TRUE,CompanySize.NORMAL,HIGH
-5,TRUE,15,10,FALSE,TRUE,TRUE,CompanySize.GIGANTISHE,HIGH
+5,TRUE,15,10,FALSE,TRUE,TRUE,CompanySize.GIGANTISHE,MEDIUM
 5,TRUE,91,94,FALSE,TRUE,TRUE,CompanySize.HUGE,HIGH
 9,TRUE,83,98,FALSE,TRUE,FALSE,CompanySize.HUGE,HIGH
 12,TRUE,58,56,FALSE,FALSE,FALSE,CompanySize.BIG,MEDIUM
@@ -219,8 +219,8 @@ DELAY,PAYED,NET-WORTH,INFLUENCE,SKARBOWKA,MEMBER,HAT,SIZE,PRIORITY
 12,FALSE,68,33,FALSE,FALSE,TRUE,CompanySize.NORMAL,LOW
 7,TRUE,12,25,FALSE,TRUE,TRUE,CompanySize.NORMAL,LOW
 8,FALSE,63,4,FALSE,FALSE,TRUE,CompanySize.GIGANTISHE,LOW
-8,TRUE,82,18,FALSE,FALSE,TRUE,CompanySize.HUGE,LOW
-5,TRUE,56,85,FALSE,FALSE,FALSE,CompanySize.HUGE,LOW
+8,TRUE,82,18,FALSE,FALSE,TRUE,CompanySize.HUGE,HIGH
+5,TRUE,56,85,FALSE,FALSE,FALSE,CompanySize.HUGE,MEDIUM
 7,FALSE,96,29,FALSE,FALSE,TRUE,CompanySize.GIGANTISHE,LOW
 9,FALSE,24,46,FALSE,TRUE,TRUE,CompanySize.BIG,LOW
 9,FALSE,38,35,FALSE,TRUE,TRUE,CompanySize.NO,LOW
diff --git a/main.py b/main.py
index 5cd0fa9..8335a3d 100644
--- a/main.py
+++ b/main.py
@@ -101,7 +101,7 @@ if __name__ == '__main__':
     label_BP.fit(['CompanySize.NO', 'CompanySize.SMALL', 'CompanySize.NORMAL', 'CompanySize.BIG', 'CompanySize.HUGE', 'CompanySize.GIGANTISHE'])
     X[:, 7] = label_BP.transform(X[:, 7])
 
-    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, train_size=0.8)
+    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.15, train_size=0.85)
 
     drugTree = DecisionTreeClassifier(criterion="entropy", max_depth=4)