Ejemplo n.º 1
0
    def test_feature_table_extraction_with_description_stream(self) -> None:
        self._init_extractor(programmatic_description_enabled=True)
        self.extractor._client.list_projects.return_value = ["default"]
        self._mock_feature_table(add_stream_source=True)

        feature_table_definition = self.extractor.extract()
        assert isinstance(feature_table_definition, TableMetadata)

        description = self.extractor.extract()
        assert isinstance(description, TableMetadata)
        expected = DescriptionMetadata(
            TestFeastExtractor._strip_margin(
                """* Created at **2020-01-01 00:00:00**
                  |"""),
            "feature_table_details",
        )
        self.assertEqual(expected.__repr__(),
                         description.description.__repr__())

        batch_source = self.extractor.extract()
        assert isinstance(batch_source, TableMetadata)
        expected = DescriptionMetadata(
            TestFeastExtractor._strip_margin("""```
                |fileOptions:
                |  fileFormat:
                |    parquetFormat: {}
                |  fileUrl: file:///some/location
                |type: BATCH_FILE
                |```"""),
            "batch_source",
        )
        self.assertEqual(expected.__repr__(),
                         batch_source.description.__repr__())

        stream_source = self.extractor.extract()
        assert isinstance(stream_source, TableMetadata)
        expected = DescriptionMetadata(
            TestFeastExtractor._strip_margin("""```
                 |createdTimestampColumn: datetime
                 |eventTimestampColumn: datetime
                 |kafkaOptions:
                 |  bootstrapServers: broker1
                 |  messageFormat:
                 |    avroFormat:
                 |      schemaJson: '{"type": "record", "name": "DriverTrips", "fields": [{"name": "driver_id",
                 |        "type": "long"}, {"name": "trips_today", "type": "int"}, {"name": "datetime",
                 |        "type": {"type": "long", "logicalType": "timestamp-micros"}}]}'
                 |  topic: driver_trips
                 |type: STREAM_KAFKA
                 |```"""),
            "stream_source",
        )
        self.assertEqual(expected.__repr__(),
                         stream_source.description.__repr__())

        self.assertIsNone(self.extractor.extract())
Ejemplo n.º 2
0
    def test_feature_table_extraction_with_description_batch(self) -> None:
        self._init_extractor(programmatic_description_enabled=True)
        self.extractor._client.list_projects.return_value = ["default"]
        self._mock_feature_table(labels={"label1": "value1"})

        feature_table_definition = self.extractor.extract()
        assert isinstance(feature_table_definition, TableMetadata)

        description = self.extractor.extract()
        assert isinstance(description, TableMetadata)
        expected = DescriptionMetadata(
            TestFeastExtractor._strip_margin(
                """* Created at **2020-01-01 00:00:00**
                  |* Labels:
                  |    * label1: **value1**
                  |"""),
            "feature_table_details",
        )
        self.assertEqual(expected.__repr__(),
                         description.description.__repr__())

        batch_source = self.extractor.extract()
        assert isinstance(batch_source, TableMetadata)
        expected = DescriptionMetadata(
            TestFeastExtractor._strip_margin("""```
                |fileOptions:
                |  fileFormat:
                |    parquetFormat: {}
                |  fileUrl: file:///some/location
                |type: BATCH_FILE
                |```"""),
            "batch_source",
        )
        self.assertEqual(expected.__repr__(),
                         batch_source.description.__repr__())

        self.assertIsNone(self.extractor.extract())