Add proto conversion test for BigQuerySource

Signed-off-by: Felix Wang <wangfelix98@gmail.com>
feast-dev · feast-ci-bot · Apr 24, 2022 · Apr 21, 2022 · Apr 23, 2022 · Apr 23, 2022
commit 99413caa5a3be66188c66241aae9d1d142e789b5
@@ -99,15 +99,9 @@ def __eq__(self, other):
             )
 
         return (
-            self.name == other.name
-            and self.bigquery_options.table == other.bigquery_options.table
-            and self.bigquery_options.query == other.bigquery_options.query
-            and self.timestamp_field == other.timestamp_field
-            and self.created_timestamp_column == other.created_timestamp_column
-            and self.field_mapping == other.field_mapping
-            and self.description == other.description
-            and self.tags == other.tags
-            and self.owner == other.owner
+            super().__eq__(other)
+            and self.table == other.table
+            and self.query == other.query
         )
 
     @property
@@ -120,7 +114,6 @@ def query(self):
 
     @staticmethod
     def from_proto(data_source: DataSourceProto):
-
         assert data_source.HasField("bigquery_options")
 
         return BigQuerySource(
@@ -144,11 +137,10 @@ def to_proto(self) -> DataSourceProto:
             description=self.description,
             tags=self.tags,
             owner=self.owner,
+            timestamp_field=self.timestamp_field,
+            created_timestamp_column=self.created_timestamp_column,
         )
 
-        data_source_proto.timestamp_field = self.timestamp_field
-        data_source_proto.created_timestamp_column = self.created_timestamp_column
-
         return data_source_proto
 
     def validate(self, config: RepoConfig):
@@ -200,42 +192,14 @@ def get_table_column_names_and_types(
 
 class BigQueryOptions:
     """
-    DataSource BigQuery options used to source features from BigQuery query
+    Configuration options for a BigQuery data source.
     """
 
     def __init__(
         self, table: Optional[str], query: Optional[str],
     ):
-        self._table = table
-        self._query = query
-
-    @property
-    def query(self):
-        """
-        Returns the BigQuery SQL query referenced by this source
-        """
-        return self._query
-
-    @query.setter
-    def query(self, query):
-        """
-        Sets the BigQuery SQL query referenced by this source
-        """
-        self._query = query
-
-    @property
-    def table(self):
-        """
-        Returns the table ref of this BQ table
-        """
-        return self._table
-
-    @table.setter
-    def table(self, table):
-        """
-        Sets the table ref of this BQ table
-        """
-        self._table = table
+        self.table = table or ""
+        self.query = query or ""
 
     @classmethod
     def from_proto(cls, bigquery_options_proto: DataSourceProto.BigQueryOptions):
@@ -248,7 +212,6 @@ def from_proto(cls, bigquery_options_proto: DataSourceProto.BigQueryOptions):
         Returns:
             Returns a BigQueryOptions object based on the bigquery_options protobuf
         """
-
         bigquery_options = cls(
             table=bigquery_options_proto.table, query=bigquery_options_proto.query,
         )
@@ -262,7 +225,6 @@ def to_proto(self) -> DataSourceProto.BigQueryOptions:
         Returns:
             BigQueryOptionsProto protobuf
         """
-
         bigquery_options_proto = DataSourceProto.BigQueryOptions(
             table=self.table, query=self.query,
         )

@@ -151,6 +151,15 @@ def test_default_data_source_kw_arg_warning():
 
 
 def test_proto_conversion():
+    bigquery_source = BigQuerySource(
+        table="test_table",
+        timestamp_field="event_timestamp",
+        created_timestamp_column="created_timestamp",
+        field_mapping={"foo": "bar"},
+        description="test description",
+        owner="test@gmail.com",
+    )
+
     file_source = FileSource(
         path="test_path",
         timestamp_field="event_timestamp",
@@ -185,6 +194,9 @@ def test_proto_conversion():
         owner="test@gmail.com",
     )
 
+    assert (
+        BigQuerySource.from_proto(bigquery_source.to_proto()) == bigquery_source_proto
+    )
     assert FileSource.from_proto(file_source.to_proto()) == file_source
     assert RedshiftSource.from_proto(redshift_source.to_proto()) == redshift_source
     assert SnowflakeSource.from_proto(snowflake_source.to_proto()) == snowflake_source