cast proto enums

Author: wei-qlu

python/lib/sift_client/_tests/resources/test_data_imports.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
+from typing import TYPE_CHECKING, cast
 
 import pytest
 from sift.common.type.v1.channel_config_pb2 import ChannelConfig as ChannelConfigProto
@@ -41,6 +42,11 @@
     TimeFormat,
 )
 
+if TYPE_CHECKING:
+    from sift.common.type.v1.channel_data_type_pb2 import (
+        ChannelDataType as ChannelDataTypeProto,
+    )
+
 
 @pytest.mark.integration
 def test_client_binding(sift_client):
@@ -421,7 +427,10 @@ def _make_flat_dataset_response(
             data_columns=[
                 ParquetDataColumnProto(
                     path=path,
-                    channel_config=ChannelConfigProto(name=path, data_type=data_type),
+                    channel_config=ChannelConfigProto(
+                        name=path,
+                        data_type=cast("ChannelDataTypeProto.ValueType", data_type),
+                    ),
                 )
                 for path, data_type in data_columns
             ],
@@ -436,7 +445,10 @@ def _make_scpr_response(time_path: str, columns: list[tuple[str, int]]) -> Parqu
             columns=[
                 ParquetColumn(
                     path=path,
-                    column_config=ChannelConfigProto(name=path, data_type=data_type),
+                    column_config=ChannelConfigProto(
+                        name=path,
+                        data_type=cast("ChannelDataTypeProto.ValueType", data_type),
+                    ),
                 )
                 for path, data_type in columns
             ],