From 98af1ef4c99ec090046dc17e61107cb70526da42 Mon Sep 17 00:00:00 2001
From: Christopher Burgess <christopher.burgess@vortexa.com>
Date: Thu, 5 Dec 2019 09:16:55 +0000
Subject: [PATCH] fix: CargoMovementsResult to_df all columns, closes #120

---
 tests/endpoints/test_cargo_movements_real.py   | 15 +++++++++++++++
 tests/endpoints/test_vessel_movements_real.py  | 14 ++++++++++++++
 vortexasdk/endpoints/cargo_movements_result.py |  5 ++++-
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/tests/endpoints/test_cargo_movements_real.py b/tests/endpoints/test_cargo_movements_real.py
index 1fa20796..b910f06b 100644
--- a/tests/endpoints/test_cargo_movements_real.py
+++ b/tests/endpoints/test_cargo_movements_real.py
@@ -46,6 +46,21 @@ def test_search_single_filter_id(self):
 
         assert len(df) == 2
 
+    def test_to_df_all_columns(self):
+        df = (
+            CargoMovements()
+            .search(
+                filter_activity="loading_state",
+                filter_products="6f11b0724c9a4e85ffa7f1445bc768f054af755a090118dcf99f14745c261653",
+                filter_time_min=datetime(2019, 8, 29),
+                filter_time_max=datetime(2019, 8, 29, 0, 10),
+            )
+            .to_df(columns="all")
+            .head(2)
+        )
+
+        assert len(df) == 2
+
     def test_search_single_filter_origin_name(self):
         df = (
             CargoMovements()
diff --git a/tests/endpoints/test_vessel_movements_real.py b/tests/endpoints/test_vessel_movements_real.py
index 64b2c90c..7a0a3760 100644
--- a/tests/endpoints/test_vessel_movements_real.py
+++ b/tests/endpoints/test_vessel_movements_real.py
@@ -15,6 +15,20 @@ def test_search(self):
 
         assert len(v) > 10
 
+    def test_to_df_all_columns(self):
+        df = (
+            VesselMovements()
+            .search(
+                filter_time_min=datetime(2017, 10, 1, 0, 0),
+                filter_time_max=datetime(2017, 10, 1, 0, 10),
+                filter_origins="rotterdam",
+            )
+            .to_df(columns="all")
+            .head(2)
+        )
+
+        assert len(df) == 2
+
     def test_search_to_dataframe(self):
         df = (
             VesselMovements()
diff --git a/vortexasdk/endpoints/cargo_movements_result.py b/vortexasdk/endpoints/cargo_movements_result.py
index 179695c8..73dccffc 100644
--- a/vortexasdk/endpoints/cargo_movements_result.py
+++ b/vortexasdk/endpoints/cargo_movements_result.py
@@ -543,7 +543,10 @@ def to_df(self, columns=None) -> pd.DataFrame:
         with Pool(os.cpu_count()) as pool:
             records = pool.map(flatten, super().to_list())
 
-        return pd.DataFrame(data=records, columns=columns)
+        if columns == "all":
+            return pd.DataFrame(data=records)
+        else:
+            return pd.DataFrame(data=records, columns=columns)
 
 
 DEFAULT_COLUMNS = [