From 6e475b90fb8a7f2f4365f288bb6bde4ffb080d8c Mon Sep 17 00:00:00 2001
From: Kelvin Muchiri <kelvinmuchiri63@gmail.com>
Date: Wed, 23 Oct 2024 14:35:34 +0300
Subject: [PATCH] refactor tests

---
 onadata/libs/tests/utils/test_csv_builder.py  | 16 ++++++++-
 .../libs/tests/utils/test_export_builder.py   | 35 +++++++++++++------
 2 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/onadata/libs/tests/utils/test_csv_builder.py b/onadata/libs/tests/utils/test_csv_builder.py
index 635b43ded1..3163ae2187 100644
--- a/onadata/libs/tests/utils/test_csv_builder.py
+++ b/onadata/libs/tests/utils/test_csv_builder.py
@@ -2115,8 +2115,22 @@ def test_extra_columns_dataview(self):
             include_images=False,
             show_choice_labels=True,
         )
+        extra_cols = [
+            "_id",
+            "_uuid",
+            "_submission_time",
+            "_date_modified",
+            "_tags",
+            "_notes",
+            "_version",
+            "_duration",
+            "_submitted_by",
+            "_total_media",
+            "_media_count",
+            "_media_all_received",
+        ]
 
-        for extra_col in csv_df_builder.extra_columns:
+        for extra_col in extra_cols:
             dataview = DataView.objects.create(
                 xform=xform,
                 name="test",
diff --git a/onadata/libs/tests/utils/test_export_builder.py b/onadata/libs/tests/utils/test_export_builder.py
index 1f656167bd..5c3b27a719 100644
--- a/onadata/libs/tests/utils/test_export_builder.py
+++ b/onadata/libs/tests/utils/test_export_builder.py
@@ -3679,13 +3679,28 @@ def test_extra_columns_dataview(self):
         )
         export_builder = ExportBuilder()
         export_builder.set_survey(self.xform.survey)
-        dataview = DataView.objects.create(
-            xform=self.xform,
-            name="test",
-            columns=["name", "_id"],
-            project=self.project,
-        )
-        fields = export_builder.get_fields(
-            dataview, export_builder.sections[0], "title"
-        )
-        self.assertEqual(fields, ["name", "_id"])
+        extra_cols = [
+            "_id",
+            "_uuid",
+            "_submission_time",
+            "_index",
+            "_parent_table_name",
+            "_parent_index",
+            "_tags",
+            "_notes",
+            "_version",
+            "_duration",
+            "_submitted_by",
+        ]
+
+        for extra_col in extra_cols:
+            dataview = DataView.objects.create(
+                xform=self.xform,
+                name="test",
+                columns=["name", extra_col],
+                project=self.project,
+            )
+            fields = export_builder.get_fields(
+                dataview, export_builder.sections[0], "title"
+            )
+            self.assertEqual(fields, ["name", extra_col])