Test RayAggregationNode functionality.
(
ray_session, ray_config, mock_context, sample_data, column_info
)
| 167 | |
| 168 | |
| 169 | def test_ray_aggregation_node( |
| 170 | ray_session, ray_config, mock_context, sample_data, column_info |
| 171 | ): |
| 172 | """Test RayAggregationNode functionality.""" |
| 173 | ray_dataset = ray.data.from_pandas(sample_data) |
| 174 | input_value = DAGValue(data=ray_dataset, format=DAGFormat.RAY) |
| 175 | dummy_node = DummyInputNode("input_node", input_value) |
| 176 | node = RayAggregationNode( |
| 177 | name="aggregation", |
| 178 | aggregations=[ |
| 179 | Aggregation(column="conv_rate", function="sum"), |
| 180 | Aggregation(column="acc_rate", function="avg"), |
| 181 | ], |
| 182 | group_by_keys=["driver_id"], |
| 183 | timestamp_col="event_timestamp", |
| 184 | config=ray_config, |
| 185 | ) |
| 186 | node.add_input(dummy_node) |
| 187 | mock_context.node_outputs = {"input_node": input_value} |
| 188 | result = node.execute(mock_context) |
| 189 | assert isinstance(result, DAGValue) |
| 190 | assert result.format == DAGFormat.RAY |
| 191 | result_df = result.data.to_pandas() |
| 192 | assert len(result_df) == 2 |
| 193 | assert "driver_id" in result_df.columns |
| 194 | assert "sum_conv_rate" in result_df.columns |
| 195 | assert "avg_acc_rate" in result_df.columns |
| 196 | |
| 197 | |
| 198 | def test_ray_join_node(ray_session, ray_config, mock_context, sample_data, column_info): |
nothing calls this directly
no test coverage detected