refactor io_tests

databrickslabs · Jul 9, 2024 · f1958b2 · f1958b2
1 parent 55e65ad
commit f1958b2
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 67 deletions.
diff --git a/python/tests/io_tests.py b/python/tests/io_tests.py
@@ -15,7 +15,7 @@ def test_write_to_delta_without_optimization_cols(self):
         table_name = "my_table_no_optimization_col"
 
         # load test data
-        input_tsdf = self.get_data_as_tsdf("input_data")
+        input_tsdf = self.get_test_df_builder("init").as_tsdf()
 
         # test write to delta
         input_tsdf.write(self.spark, table_name)
@@ -30,7 +30,7 @@ def test_write_to_delta_with_optimization_cols(self):
         table_name = "my_table_optimization_col"
 
         # load test data
-        input_tsdf = self.get_data_as_tsdf("input_data")
+        input_tsdf = self.get_test_df_builder("init").as_tsdf()
 
         # test write to delta
         input_tsdf.write(self.spark, table_name, ["date"])
@@ -45,7 +45,7 @@ def test_write_to_delta_bad_dbr_environment_logging(self):
         table_name = "my_table_optimization_col_fails"
 
         # load test data
-        input_tsdf = self.get_data_as_tsdf("input_data")
+        input_tsdf = self.get_test_df_builder("init").as_tsdf()
 
         if pkg_version.parse(DELTA_VERSION) < pkg_version.parse("2.0.0"):
 

diff --git a/python/tests/unit_test_data/io_tests.json b/python/tests/unit_test_data/io_tests.json
@@ -1,83 +1,85 @@
 {
   "__SharedData": {
-    "input_data": {
-      "schema": "symbol string, date string, event_ts string, trade_pr float, trade_pr_2 float",
-      "ts_col": "event_ts",
-      "partition_cols": [
-        "symbol"
-      ],
-      "data": [
-        [
-          "S1",
-          "SAME_DT",
-          "2020-08-01 00:00:10",
-          349.21,
-          10.0
-        ],
-        [
-          "S1",
-          "SAME_DT",
-          "2020-08-01 00:00:11",
-          340.21,
-          9.0
-        ],
-        [
-          "S1",
-          "SAME_DT",
-          "2020-08-01 00:01:12",
-          353.32,
-          8.0
-        ],
-        [
-          "S1",
-          "SAME_DT",
-          "2020-08-01 00:01:13",
-          351.32,
-          7.0
-        ],
-        [
-          "S1",
-          "SAME_DT",
-          "2020-08-01 00:01:14",
-          350.32,
-          6.0
-        ],
-        [
-          "S1",
-          "SAME_DT",
-          "2020-09-01 00:01:12",
-          361.1,
-          5.0
-        ],
-        [
-          "S1",
-          "SAME_DT",
-          "2020-09-01 00:19:12",
-          362.1,
-          4.0
+    "init": {
+      "tsdf": {
+        "ts_col": "event_ts",
+        "partition_cols": ["symbol"]
+      },
+      "df": {
+        "schema": "symbol string, date string, event_ts string, trade_pr float, trade_pr_2 float",
+        "data": [
+          [
+            "S1",
+            "SAME_DT",
+            "2020-08-01 00:00:10",
+            349.21,
+            10.0
+          ],
+          [
+            "S1",
+            "SAME_DT",
+            "2020-08-01 00:00:11",
+            340.21,
+            9.0
+          ],
+          [
+            "S1",
+            "SAME_DT",
+            "2020-08-01 00:01:12",
+            353.32,
+            8.0
+          ],
+          [
+            "S1",
+            "SAME_DT",
+            "2020-08-01 00:01:13",
+            351.32,
+            7.0
+          ],
+          [
+            "S1",
+            "SAME_DT",
+            "2020-08-01 00:01:14",
+            350.32,
+            6.0
+          ],
+          [
+            "S1",
+            "SAME_DT",
+            "2020-09-01 00:01:12",
+            361.1,
+            5.0
+          ],
+          [
+            "S1",
+            "SAME_DT",
+            "2020-09-01 00:19:12",
+            362.1,
+            4.0
+          ]
         ]
-      ]
+      }
     }
   },
   "DeltaWriteTest": {
     "test_write_to_delta_without_optimization_cols": {
-      "input_data": {
-        "$ref": "#/__SharedData/input_data"
+      "init": {
+        "$ref": "#/__SharedData/init"
       }
     },
     "test_write_to_delta_with_optimization_cols": {
-      "input_data": {
-        "$ref": "#/__SharedData/input_data"
+      "init": {
+        "$ref": "#/__SharedData/init"
       }
     },
     "test_write_to_delta_non_dbr_environment_logging": {
-      "input_data": {
-        "$ref": "#/__SharedData/input_data"
+      "init": {
+        "$ref": "#/__SharedData/init"
       }
     },
     "test_write_to_delta_bad_dbr_environment_logging": {
-      "input_data": {
-        "$ref": "#/__SharedData/input_data"
+      "init": {
+        "$ref": "#/__SharedData/init"
       }
     }
   }