additional checks for idf dataframe equality

python/tests/base.py

@@ -158,9 +158,9 @@ def as_idf(self) -> IntervalsDF:
         """
         sdf = self.as_sdf()
         if self.idf_construct is not None:
-            return getattr(IntervalsDF, self.idf_construct)(sdf, **self.tsdf)
+            return getattr(IntervalsDF, self.idf_construct)(sdf, **self.idf)
         else:
-            return IntervalsDF(self.as_sdf(), **self.tsdf)
+            return IntervalsDF(self.as_sdf(), **self.idf)
 
 
 class SparkTest(unittest.TestCase):
@@ -305,8 +305,8 @@ def assertSchemaContainsField(self, schema, field):
 
     def assertDataFrameEquality(
         self,
-        df1: Union[TSDF, DataFrame],
-        df2: Union[TSDF, DataFrame],
+        df1: Union[TSDF, DataFrame, IntervalsDF],
+        df2: Union[TSDF, DataFrame, IntervalsDF],
         ignore_row_order: bool = False,
         ignore_column_order: bool = True,
         ignore_nullable: bool = True,
@@ -324,6 +324,14 @@ def assertDataFrameEquality(
             df1 = df1.df
             df2 = df2.df
 
+        # Handle IDFs
+        if isinstance(df1, IntervalsDF):
+            # df2 must also be a IntervalsDF
+            self.assertIsInstance(df2, IntervalsDF)
+            # get the underlying Spark DataFrames
+            df1 = df1.df
+            df2 = df2.df
+
         # handle DataFrames
         assert_df_equality(
             df1,