add: test function for convert_data_webdataset

Author: yztxwd

tests/test_writer_loader.py

@@ -3,7 +3,7 @@
 sys.path.insert(0, "./")
 import pandas as pd
 from seqchromloader import SeqChromDatasetByDataFrame, SeqChromDatasetByBed, SeqChromDatasetByWds, SeqChromDataModule
-from seqchromloader import dump_data_webdataset
+from seqchromloader import dump_data_webdataset, convert_data_webdataset
 
 import unittest
 import tempfile
@@ -85,6 +85,20 @@ def test_wds_loader_transform(self):
         self.assertEqual(target[0].item(), 0.0)
         self.assertEqual(label[1].item(), 1)
 
+    def test_wds_convert_loader(self):
+        convert_data_webdataset("data/test_0.tar.gz", "test_0_convert.tar.gz",
+                                        transforms={"seq": test_seq_transform,
+                                                    "chrom": test_chrom_transform,
+                                                    "target": test_target_transform})
+        it = iter(SeqChromDatasetByWds(["test_0_convert.tar.gz"],
+                  dataloader_kws={"batch_size":3}))
+        seq, chrom, target, label = next(it)
+
+        self.assertEqual(seq[0,0,3].item(), 2.0)
+        self.assertAlmostEqual(chrom[0,0,3].item(), 4.0/3)
+        self.assertEqual(target[0].item(), 0.0)
+        self.assertEqual(label[1].item(), 1)
+    
     def test_df_loader(self):
         dataframe = pd.read_table("data/sample.bed", header=None, sep="\t", names=['chrom', 'start', 'end', 'label', 'score', 'strand' ])
         it = iter(SeqChromDatasetByDataFrame(