From db2cdce98a5ab63535592791cb9be18c21188175 Mon Sep 17 00:00:00 2001
From: migueluzcategui <miguel.uzcategui@vortexa.com>
Date: Tue, 5 Sep 2023 14:45:14 +0100
Subject: [PATCH] feat: improve tests

---
 tests/api/test_cargo_movement.py              | 18 ++++++++
 tests/endpoints/test_vessels_search_result.py | 46 +++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/tests/api/test_cargo_movement.py b/tests/api/test_cargo_movement.py
index d5354d43..2361fe19 100644
--- a/tests/api/test_cargo_movement.py
+++ b/tests/api/test_cargo_movement.py
@@ -24,6 +24,12 @@ class TestCargoMovement(TestCase):
                 "dwt": 5260,
                 "cubic_capacity": 6100,
                 "vessel_class": "tiny_tanker",
+                "classes": [
+                    {"id": "vlcc_plus", "label": "VLCC+", "layer": "legacy"},
+                    {"id": "oil", "label": "Oil Tankers", "layer": "group"},
+                    {"id": "oil_vlcc", "label": "VLCC+", "layer": "coarse"},
+                    {"id": "oil_vlcc", "label": "VLCC+", "layer": "granular"},
+                ],
                 "corporate_entities": [
                     {
                         "id": "f9bd45e65e292909a7b751b0026dcf7795c6194b3c0712910a241caee32c99b8",
@@ -156,6 +162,18 @@ def test_convert_to_flat_dict(self):
             "vessels.0.start_timestamp": "2019-10-18T21:38:34+0000",
             "vessels.0.status": "vessel_status_laden_known",
             "vessels.0.vessel_class": "tiny_tanker",
+            "vessels.0.classes.0.id": "vlcc_plus",
+            "vessels.0.classes.0.layer": "legacy",
+            "vessels.0.classes.0.label": "VLCC+",
+            "vessels.0.classes.1.id": "oil",
+            "vessels.0.classes.1.layer": "group",
+            "vessels.0.classes.1.label": "Oil Tankers",
+            "vessels.0.classes.2.id": "oil_vlcc",
+            "vessels.0.classes.2.layer": "coarse",
+            "vessels.0.classes.2.label": "VLCC+",
+            "vessels.0.classes.3.id": "oil_vlcc",
+            "vessels.0.classes.3.layer": "granular",
+            "vessels.0.classes.3.label": "VLCC+",
             "vessels.0.scrubber": None,
             "vessels.0.flag": None,
             "vessels.0.voyage_id": "401f0e74fc42401248a484aca2e9955dea885378796f7f4d0bc8e92c35ea270a",
diff --git a/tests/endpoints/test_vessels_search_result.py b/tests/endpoints/test_vessels_search_result.py
index 68c9ed75..33ca1f88 100644
--- a/tests/endpoints/test_vessels_search_result.py
+++ b/tests/endpoints/test_vessels_search_result.py
@@ -105,6 +105,52 @@ def test_to_df(self):
                 "corporate_entities": {0: [], 1: []},
                 "leaf": {0: True, 1: True},
                 "vessel_class": {0: "tiny_tanker", 1: "tiny_tanker"},
+                "classes": {
+                    0: [
+                        {
+                            "id": "vlcc_plus",
+                            "label": "VLCC+",
+                            "layer": "legacy",
+                        },
+                        {
+                            "id": "oil",
+                            "label": "Oil Tankers",
+                            "layer": "group",
+                        },
+                        {
+                            "id": "oil_vlcc",
+                            "label": "VLCC+",
+                            "layer": "coarse",
+                        },
+                        {
+                            "id": "oil_vlcc",
+                            "label": "VLCC+",
+                            "layer": "granular",
+                        },
+                    ],
+                    1: [
+                        {
+                            "id": "vlcc_plus",
+                            "label": "VLCC+",
+                            "layer": "legacy",
+                        },
+                        {
+                            "id": "oil",
+                            "label": "Oil Tankers",
+                            "layer": "group",
+                        },
+                        {
+                            "id": "oil_vlcc",
+                            "label": "VLCC+",
+                            "layer": "coarse",
+                        },
+                        {
+                            "id": "oil_vlcc",
+                            "label": "VLCC+",
+                            "layer": "granular",
+                        },
+                    ],
+                },
                 "cubic_capacity": {0: 7060, 1: 6318},
                 "tags": {
                     0: [