Issue #782 unit tests for function calculate_vwap_twap with 2 differe…

…nt resampling rules

im_v2/common/data/transform/test/test_transform_utils.py

@@ -3,6 +3,7 @@
 
 import pandas as pd
 import pytest
+import random
 
 import data_schema.dataset_schema_utils as dsdascut
 import helpers.hdbg as hdbg
@@ -14,6 +15,7 @@
 import im_v2.common.data.extract.extract_utils as imvcdeexut
 import im_v2.common.data.transform.resample_daily_bid_ask_data as imvcdtrdbad
 import im_v2.common.data.transform.transform_utils as imvcdttrut
+import research_amp.transform as ramptran
 
 
 class TestGetVendorEpochUnit(hunitest.TestCase):
@@ -781,20 +783,85 @@ def _get_test_data(self) -> None:
 # #############################################################################
 
 
-class TestCalculateVwap(hunitest.TestCase):
+class TestCalculateVwapTwap(hunitest.TestCase):
 
-    def generate_test_data(self):
+    def get_test_data(self):
         """
         Fetch data for price and volume columns with timestamp index
-        returns: pandas dataframe with timestamp index, price column and volume column
+        returns: pandas dataframe with timestamp index, close and volume column with asset_id as full_sumbol
         """
+        #
+        index_timestamp = pd.date_range("2024-01-01", periods=2, freq="min")
+        close = [34501.279, 29952.436]
+        volume = [229894.0, 99837.4]
+        asset_id = [18109, 15479]
+        d = {"timestamp":index_timestamp, "close":close, "volume":volume, "full_symbol":asset_id}
+        df = pd.DataFrame(data=d).set_index("timestamp")
+        expected_df_columns = pd.MultiIndex.from_product(
+            [["close", "twap", "volume", "vwap"], sorted(asset_id)]
+        )
+        return df, expected_df_columns
 
-    def test_calculate_vwap_with_resample_kwargs(self) -> None:
+    def test_calculate_vwap_twap_with_5T_resample_kwargs(self) -> None:
         """
-        Verify calculated price for given volume from data WITH resample rules
+        Verify calculated price for given volume from data WITH resample rule
         """
+        df, expected_df_columns = self.get_test_data()
+        result_df = ramptran.calculate_vwap_twap(
+            df, "5T"
+        )
+        #Expected Values
+        expected_length = 2
+        expected_column_unique_values = None
+        expected_signature=r"""
+        # df=
+        index=[2024-01-01 00:00:00, 2024-01-01 00:05:00]
+        columns=('close', 15479),('close', 18109),('twap', 15479),('twap', 18109),('volume', 15479),('volume', 18109),('vwap', 15479),('vwap', 18109)
+        shape=(2, 8)
+                            close	        twap	        volume	        vwap
+                            15479	18109	15479	18109	15479	18109	15479	18109
+                  timestamp								
+        2024-01-01 00:00:00	NaN	34501.279	NaN	34501.279	NaN	229894.0	NaN	34501.279
+        2024-01-01 00:05:00	29952.436 NaN	29952.436 NaN	99837.4	NaN	    29952.436 NaN
+        """
+        # Check.
+        self.check_df_output(
+            result_df,
+            expected_length,
+            expected_df_columns,
+            expected_column_unique_values,
+            expected_signature,
+        )
+
 
-    def test_calculate_vwap_no_resample_kwargs(self) -> None:
+    def test_calculate_vwap_twap_with_halfT_resample_kwargs(self) -> None:
         """
         Verify calculated price for given volume from data WITHOUT resample rules
-        """
+        """
+        df, expected_df_columns = self.get_test_data()
+        result_df = ramptran.calculate_vwap_twap(
+            df, "0.5T"
+        )
+        #Expected Values
+        expected_length = 3
+        expected_column_unique_values = None
+        expected_signature = r"""
+        # df=
+        index=[2024-01-01 00:00:00, 2024-01-01 00:01:00]
+        columns=('close', 15479),('close', 18109),('twap', 15479),('twap', 18109),('volume', 15479),('volume', 18109),('vwap', 15479),('vwap', 18109)
+        shape=(3, 8)
+                            close	        twap	        volume	        vwap
+                            15479	18109	15479	18109	15479	18109	15479	18109
+                  timestamp								
+        2024-01-01 00:00:00	NaN	34501.279	NaN	34501.279	NaN	229894.0	NaN	34501.279
+        2024-01-01 00:00:30	NaN	      NaN 	NaN 	  NaN	NaN	     NaN	NaN	      NaN
+        2024-01-01 00:01:00	29952.436 NaN	29952.436 NaN	99837.4	 NaN	29952.436 NaN
+        """
+        # Check.
+        self.check_df_output(
+            result_df,
+            expected_length,
+            expected_df_columns,
+            expected_column_unique_values,
+            expected_signature,
+        )