Sorr task918 add unit test for split universe

im_v2/common/data/extract/test/test_extract_utils.py

@@ -826,6 +826,62 @@ def test_download_and_resample_bid_ask_data(self) -> None:
     not henv.execute_repo_config_code("is_CK_S3_available()"),
     reason="Run only if CK S3 is available",
 )
+class TestSplitUniverse(hunitest.TestCase):
+    # A universe list from the example.
+    universe = [
+        "ALICE_USDT",
+        "GALA_USDT",
+        "FLOW_USDT",
+        "HBAR_USDT",
+        "INJ_USDT",
+        "NEAR_USDT",
+    ]
+
+    def test1(self) -> None:
+        """
+        Test if both group_size and universe_part are 0.
+        """
+        universe_subset = imvcdeexut._split_universe(self.universe, 0, 0)
+        actual_str = " ".join([str(element) for element in universe_subset])
+        expected_str = r"""
+
+        """
+        self.assert_equal(actual_str, expected_str, fuzzy_match=True)
+
+    def test2(self) -> None:
+        """
+        Test if group_size is 6 and universe_part is 1.
+
+        (This returns the whole list in the 1st part)
+        """
+        universe_subset = imvcdeexut._split_universe(self.universe, 6, 1)
+        actual_str = " ".join([str(element) for element in universe_subset])
+        expected_str = r"""
+        ALICE_USDT GALA_USDT FLOW_USDT HBAR_USDT INJ_USDT NEAR_USDT
+        """
+        self.assert_equal(actual_str, expected_str, fuzzy_match=True)
+
+    def test3(self) -> None:
+        """
+        Test if roup_size is 6 and universe_part is 2.
+
+        (This returns an empty list, since no groups in the part)
+        """
+        universe_subset = imvcdeexut._split_universe(self.universe, 6, 2)
+        actual_str = " ".join([str(element) for element in universe_subset])
+        expected_str = r"""
+
+        """
+        self.assert_equal(actual_str, expected_str, fuzzy_match=True)
+
+    def test4(self) -> None:
+        """
+        Test to check for error.
+        """
+        with pytest.raises(RuntimeError):
+            (imvcdeexut._split_universe(self.universe, 6, 6))
+
+
 class TestDownloadHistoricalData1(hmoto.S3Mock_TestCase):
     def call_download_historical_data(
         self, incremental: bool, *, assert_on_missing_data: bool = False