diff --git a/soitool/enumerates.py b/soitool/enumerates.py
index 14d9938b6a3c48c01cd4f3d41b41f914d5ec0eb2..30940b11503d86b51a70c16d701c11d6b78980d6 100644
--- a/soitool/enumerates.py
+++ b/soitool/enumerates.py
@@ -13,4 +13,4 @@ class CodebookSort(Enum):
     """Enumerate with ways to sort codebook."""
 
     WORD = 0
-    CODE = 1
\ No newline at end of file
+    CODE = 1
diff --git a/test/test_database.py b/test/test_database.py
index 6609df817bbc514db03a1f34b3c74e14fce04a5d..e43cdc2e2f0496211d46aa03a34ad3552ce6d487 100644
--- a/test/test_database.py
+++ b/test/test_database.py
@@ -3,6 +3,7 @@ import os
 from pathlib import Path
 import unittest
 import json
+from soitool.enumerates import CodebookSize
 from soitool.database import Database
 from soitool.coder import get_code_length_needed
 
@@ -102,8 +103,8 @@ class DatabaseTest(unittest.TestCase):
         for i, entry in enumerate(expected):
             self.assertEqual(entry["word"], actual[i][0])
             self.assertEqual(entry["category"], actual[i][1])
-            self.assertRegex(actual[i][3], "[A-Z]{" + str(code_len) + "}")
             self.assertEqual(entry["type"], actual[i][2])
+            self.assertRegex(actual[i][3], "[A-Z]{" + str(code_len) + "}")
 
     def test_get_categories(self):
         """Assert function get_categories works as expected."""
@@ -161,7 +162,7 @@ class DatabaseTest(unittest.TestCase):
             if entry["type"] == 1:
                 expected.append(entry)
         # Get small codebook from db
-        actual = self.database.get_codebook(small=True)
+        actual = self.database.get_codebook(codebook_size=CodebookSize.FULL)
         # Compare lenght
         self.assertEqual(len(expected), len(actual))
         # Get expected length of codes