diff --git a/andromeda/tests/test_dataset.py b/andromeda/tests/test_dataset.py
index 14b31e1..035d85b 100644
--- a/andromeda/tests/test_dataset.py
+++ b/andromeda/tests/test_dataset.py
@@ -53,7 +53,7 @@ def fake_get_path():
 
         self.assertEqual(weights.keys(), set(["R1", "B1", "G1"]))
         self.assertEqual(len(image_coordinates), 3)
-        self.assertEqual(image_coordinates[0].keys(), set(["x", "y", "label", "url"]))
+        self.assertEqual(image_coordinates[0].keys(), set(["x", "y", "label", "url", "R1", "B1", "G1"]))
         self.assertEqual(image_coordinates[0]["label"], "p1")
         self.assertEqual(image_coordinates[0]["url"], "https://example.com/p1.jpg")
         self.assertEqual(image_coordinates[1]["url"], "")