From 30728e4247d4dd161c8127fa4a38e70673986164 Mon Sep 17 00:00:00 2001 From: Yu Ishikawa Date: Tue, 14 Dec 2021 11:39:45 +0900 Subject: [PATCH] Support dbt 1.0 02 (#31) * Restructure how to generate models * Update * Update * Test on Python 3.8 and 3.9 * Don't support 3.9 yet --- .github/workflows/checks.yml | 9 +- Makefile | 5 +- dbt_artifacts_loader/api/rest_api_v2.py | 12 +- .../dbt/base_bigquery_model.py | 10 +- dbt_artifacts_loader/dbt/model_factory.py | 59 --------- dbt_artifacts_loader/dbt/utils.py | 107 +++------------- dbt_artifacts_loader/dbt/version_map.py | 114 ++++++++++++++++++ terraform/example/main.tf | 2 +- tests/dbt/test_model_factory.py | 45 ------- tests/dbt/test_utils.py | 39 +++++- 10 files changed, 197 insertions(+), 205 deletions(-) delete mode 100644 dbt_artifacts_loader/dbt/model_factory.py create mode 100644 dbt_artifacts_loader/dbt/version_map.py delete mode 100644 tests/dbt/test_model_factory.py diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 2a093bc..532fcc5 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -6,6 +6,11 @@ on: jobs: python: + strategy: + matrix: + include: + - python_version: 3.8 + runs-on: ubuntu-20.04 env: ENV_FILE: ".env/.env.test" @@ -14,14 +19,14 @@ jobs: uses: actions/checkout@v2 - uses: actions/setup-python@v2 with: - python-version: 3.8 + python-version: ${{ matrix.python_version }} - name: install dependencies run: | bash ci/setup.sh - name: lint run: | bash ci/lint_python.sh - - name: + - name: test run: | export ENV_FILE="${{ env.ENV_FILE }}" bash ci/run_python_tests.sh diff --git a/Makefile b/Makefile index 99a4ead..9c28cc1 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,6 @@ -DOCKER_IMAGE_BASE = gcr.io/ubie-yu-sandbox/dbt-artifacts-loader -TAG = "v1.1.0-rc3" +PROJECT_ID = YOUR-GCP-PROJECT +TAG = v1.2.0-dev2 +DOCKER_IMAGE_BASE = gcr.io/$(PROJECT_ID)/dbt-artifacts-loader .PHONEY: setup diff --git a/dbt_artifacts_loader/api/rest_api_v2.py b/dbt_artifacts_loader/api/rest_api_v2.py index 2e24912..35b4819 100644 --- a/dbt_artifacts_loader/api/rest_api_v2.py +++ b/dbt_artifacts_loader/api/rest_api_v2.py @@ -28,12 +28,16 @@ from google.cloud import bigquery -from dbt_artifacts_loader.dbt.model_factory import get_model_class +from dbt_artifacts_loader.dbt.version_map import ArtifactsTypes from dbt_artifacts_loader.utils import download_gcs_object_as_text +from dbt_artifacts_loader.dbt.utils import ( + get_dbt_schema_version, + get_default_load_job_config, + load_table_from_json, + get_artifact_type_by_id, + get_destination_table, + get_model_class) from dbt_artifacts_loader.api import config -from dbt_artifacts_loader.dbt.utils import get_dbt_schema_version, get_default_load_job_config, load_table_from_json, \ - get_artifact_type_by_id, get_destination_table -from dbt_artifacts_loader.dbt.utils import ArtifactsTypes app = FastAPI() diff --git a/dbt_artifacts_loader/dbt/base_bigquery_model.py b/dbt_artifacts_loader/dbt/base_bigquery_model.py index 95e4a4b..ce2465d 100644 --- a/dbt_artifacts_loader/dbt/base_bigquery_model.py +++ b/dbt_artifacts_loader/dbt/base_bigquery_model.py @@ -18,7 +18,8 @@ import json from enum import Enum from typing import Dict, Any, Optional, Union, List -from datetime import datetime, date +import datetime +from datetime import date, datetime import inspect # pylint: disable=E0611 @@ -28,7 +29,12 @@ from google.cloud import bigquery -from dbt_artifacts_loader.dbt.utils import datetime_handler + +def datetime_handler(x): + """The handler is used to deal with date and datetime""" + if isinstance(x, (datetime.datetime, datetime.date, datetime, date)): + return x.isoformat() + raise TypeError(f'Type {type(x)} not serializable') class TypingUtils: diff --git a/dbt_artifacts_loader/dbt/model_factory.py b/dbt_artifacts_loader/dbt/model_factory.py deleted file mode 100644 index 120de4d..0000000 --- a/dbt_artifacts_loader/dbt/model_factory.py +++ /dev/null @@ -1,59 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -from dbt_artifacts_loader.dbt.utils import ArtifactsTypes -from dbt_artifacts_loader.dbt.v1.catalog import CatalogV1 -from dbt_artifacts_loader.dbt.v1.manifest import ManifestV1 -from dbt_artifacts_loader.dbt.v1.run_results import RunResultsV1 -from dbt_artifacts_loader.dbt.v1.sources import SourcesV1 -from dbt_artifacts_loader.dbt.v2.manifest import ManifestV2 -from dbt_artifacts_loader.dbt.v2.run_results import RunResultsV2 -from dbt_artifacts_loader.dbt.v2.sources import SourcesV2 -from dbt_artifacts_loader.dbt.v3.manifest import ManifestV3 -from dbt_artifacts_loader.dbt.v3.run_results import RunResultsV3 - - -def get_model_class(artifact_type: ArtifactsTypes): - """Get the model class - - Args: - artifact_type (ArtifactsTypes): artifact type - - Returns: - the model class - """ - # v1 - if artifact_type == ArtifactsTypes.CATALOG_V1: - return CatalogV1 - elif artifact_type == ArtifactsTypes.MANIFEST_V1: - return ManifestV1 - elif artifact_type == ArtifactsTypes.RUN_RESULTS_V1: - return RunResultsV1 - elif artifact_type == ArtifactsTypes.SOURCES_V1: - return SourcesV1 - # v2 - elif artifact_type == ArtifactsTypes.MANIFEST_V2: - return ManifestV2 - elif artifact_type == ArtifactsTypes.RUN_RESULTS_V2: - return RunResultsV2 - elif artifact_type == ArtifactsTypes.SOURCES_V2: - return SourcesV2 - # v3 - elif artifact_type == ArtifactsTypes.MANIFEST_V3: - return ManifestV3 - elif artifact_type == ArtifactsTypes.RUN_RESULTS_V3: - return RunResultsV3 - raise ValueError(f"No such an artifact {artifact_type}") diff --git a/dbt_artifacts_loader/dbt/utils.py b/dbt_artifacts_loader/dbt/utils.py index 916d831..741aba2 100644 --- a/dbt_artifacts_loader/dbt/utils.py +++ b/dbt_artifacts_loader/dbt/utils.py @@ -15,97 +15,12 @@ # limitations under the License. # # -from enum import Enum -import datetime -from datetime import date, datetime -from typing import Optional, List -from dataclasses import dataclass +from typing import Optional, List, Type from google.cloud import bigquery - -class DestinationTables(Enum): - # V1 - CATALOG_V1 = "catalog_v1" - MANIFEST_V1 = "manifest_v1" - RUN_RESULTS_V1 = "run_results_v1" - SOURCES_V1 = "sources_v1" - # V2 - MANIFEST_V2 = "manifest_v2" - RUN_RESULTS_V2 = "run_results_v2" - SOURCES_V2 = "sources_v2" - # V3 - MANIFEST_V3 = "manifest_v3" - RUN_RESULTS_V3 = "run_results_v3" - SOURCES_V3 = "sources_v3" - # V4 - MANIFEST_V4 = "manifest_v4" - RUN_RESULTS_V4 = "run_results_v4" - - -class ArtifactsTypes(Enum): - # V1 - CATALOG_V1 = "CatalogV1" - MANIFEST_V1 = "ManifestV1" - RUN_RESULTS_V1 = "RunResultsV1" - SOURCES_V1 = "SourcesV1" - # V2 - MANIFEST_V2 = "ManifestV2" - RUN_RESULTS_V2 = "RunResultsV2" - SOURCES_V2 = "SourcesV2" - # V3 - MANIFEST_V3 = "ManifestV3" - RUN_RESULTS_V3 = "RunResultsV3" - SOURCES_V3 = "SourcesV3" - # V4 - MANIFEST_V4 = "ManifestV4" - RUN_RESULTS_V4 = "RunResultsV4" - - -@dataclass -class ArtifactInfo: - dbt_schema_version: str - artifact_type: ArtifactsTypes - destination_table: DestinationTables - - -ARTIFACT_INFO = { - # V1 - "CATALOG_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/catalog/v1.json", - ArtifactsTypes.CATALOG_V1, DestinationTables.CATALOG_V1), - "MANIFEST_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v1.json", - ArtifactsTypes.MANIFEST_V1, DestinationTables.MANIFEST_V1), - "RUN_RESULTS_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v1.json", - ArtifactsTypes.RUN_RESULTS_V1, DestinationTables.RUN_RESULTS_V1), - "SOURCES_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/sources/v1.json", - ArtifactsTypes.SOURCES_V1, DestinationTables.SOURCES_V1), - # V2 - "MANIFEST_V2": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v2.json", - ArtifactsTypes.MANIFEST_V2, DestinationTables.MANIFEST_V2), - "RUN_RESULTS_V2": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v2.json", - ArtifactsTypes.RUN_RESULTS_V2, DestinationTables.RUN_RESULTS_V2), - "SOURCES_V2": ArtifactInfo("https://schemas.getdbt.com/dbt/sources/v2.json", - ArtifactsTypes.SOURCES_V2, DestinationTables.SOURCES_V2), - # V3 - "MANIFEST_V3": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v3.json", - ArtifactsTypes.MANIFEST_V3, DestinationTables.MANIFEST_V3), - "RUN_RESULTS_V3": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v3.json", - ArtifactsTypes.RUN_RESULTS_V3, DestinationTables.RUN_RESULTS_V3), - "SOURCES_V3": ArtifactInfo("https://schemas.getdbt.com/dbt/sources/v3.json", - ArtifactsTypes.SOURCES_V3, DestinationTables.SOURCES_V3), - # V4 - "MANIFEST_V4": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v4.json", - ArtifactsTypes.MANIFEST_V4, DestinationTables.MANIFEST_V4), - "RUN_RESULTS_V4": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v4.json", - ArtifactsTypes.RUN_RESULTS_V4, DestinationTables.RUN_RESULTS_V4), -} - - -def datetime_handler(x): - """The handler is used to deal with date and datetime""" - if isinstance(x, (datetime.datetime, datetime.date, datetime, date)): - return x.isoformat() - raise TypeError(f'Type {type(x)} not serializable') +from dbt_artifacts_loader.dbt.base_bigquery_model import BaseBigQueryModel +from dbt_artifacts_loader.dbt.version_map import ArtifactsTypes, ARTIFACT_INFO, DestinationTables def get_dbt_schema_version(artifact_json: dict) -> str: @@ -178,3 +93,19 @@ def load_table_from_json( # pylint: disable=W0703 except Exception as e: raise RuntimeError(job.errors) from e + + +def get_model_class(artifact_type: ArtifactsTypes) -> Type[BaseBigQueryModel]: + """Get the model class + + Args: + artifact_type (ArtifactsTypes): artifact type + + Returns: + the model class + """ + # v1 + for _, artifact_info in ARTIFACT_INFO.items(): + if artifact_type == artifact_info.artifact_type: + return artifact_info.model_class + raise ValueError(f"No such an artifact {artifact_type}") diff --git a/dbt_artifacts_loader/dbt/version_map.py b/dbt_artifacts_loader/dbt/version_map.py new file mode 100644 index 0000000..bab636b --- /dev/null +++ b/dbt_artifacts_loader/dbt/version_map.py @@ -0,0 +1,114 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from typing import Type +from enum import Enum +from dataclasses import dataclass + +from dbt_artifacts_loader.dbt.base_bigquery_model import BaseBigQueryModel + +from dbt_artifacts_loader.dbt.v1.catalog import CatalogV1 +from dbt_artifacts_loader.dbt.v1.manifest import ManifestV1 +from dbt_artifacts_loader.dbt.v1.run_results import RunResultsV1 +from dbt_artifacts_loader.dbt.v1.sources import SourcesV1 +from dbt_artifacts_loader.dbt.v2.manifest import ManifestV2 +from dbt_artifacts_loader.dbt.v2.run_results import RunResultsV2 +from dbt_artifacts_loader.dbt.v2.sources import SourcesV2 +from dbt_artifacts_loader.dbt.v3.manifest import ManifestV3 +from dbt_artifacts_loader.dbt.v3.run_results import RunResultsV3 +from dbt_artifacts_loader.dbt.v3.sources import SourcesV3 +from dbt_artifacts_loader.dbt.v4.manifest import ManifestV4 +from dbt_artifacts_loader.dbt.v4.run_results import RunResultsV4 + + +class DestinationTables(Enum): + """Destination tables""" + # V1 + CATALOG_V1 = "catalog_v1" + MANIFEST_V1 = "manifest_v1" + RUN_RESULTS_V1 = "run_results_v1" + SOURCES_V1 = "sources_v1" + # V2 + MANIFEST_V2 = "manifest_v2" + RUN_RESULTS_V2 = "run_results_v2" + SOURCES_V2 = "sources_v2" + # V3 + MANIFEST_V3 = "manifest_v3" + RUN_RESULTS_V3 = "run_results_v3" + SOURCES_V3 = "sources_v3" + # V4 + MANIFEST_V4 = "manifest_v4" + RUN_RESULTS_V4 = "run_results_v4" + + +class ArtifactsTypes(Enum): + """Dbt artifacts types""" + # V1 + CATALOG_V1 = "CatalogV1" + MANIFEST_V1 = "ManifestV1" + RUN_RESULTS_V1 = "RunResultsV1" + SOURCES_V1 = "SourcesV1" + # V2 + MANIFEST_V2 = "ManifestV2" + RUN_RESULTS_V2 = "RunResultsV2" + SOURCES_V2 = "SourcesV2" + # V3 + MANIFEST_V3 = "ManifestV3" + RUN_RESULTS_V3 = "RunResultsV3" + SOURCES_V3 = "SourcesV3" + # V4 + MANIFEST_V4 = "ManifestV4" + RUN_RESULTS_V4 = "RunResultsV4" + + +@dataclass +class ArtifactInfo: + dbt_schema_version: str + artifact_type: ArtifactsTypes + destination_table: DestinationTables + model_class: Type[BaseBigQueryModel] + + +ARTIFACT_INFO = { + # V1 + "CATALOG_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/catalog/v1.json", + ArtifactsTypes.CATALOG_V1, DestinationTables.CATALOG_V1, CatalogV1), + "MANIFEST_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v1.json", + ArtifactsTypes.MANIFEST_V1, DestinationTables.MANIFEST_V1, ManifestV1), + "RUN_RESULTS_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v1.json", + ArtifactsTypes.RUN_RESULTS_V1, DestinationTables.RUN_RESULTS_V1, RunResultsV1), + "SOURCES_V1": ArtifactInfo("https://schemas.getdbt.com/dbt/sources/v1.json", + ArtifactsTypes.SOURCES_V1, DestinationTables.SOURCES_V1, SourcesV1), + # V2 + "MANIFEST_V2": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v2.json", + ArtifactsTypes.MANIFEST_V2, DestinationTables.MANIFEST_V2, ManifestV2), + "RUN_RESULTS_V2": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v2.json", + ArtifactsTypes.RUN_RESULTS_V2, DestinationTables.RUN_RESULTS_V2, RunResultsV2), + "SOURCES_V2": ArtifactInfo("https://schemas.getdbt.com/dbt/sources/v2.json", + ArtifactsTypes.SOURCES_V2, DestinationTables.SOURCES_V2, SourcesV2), + # V3 + "MANIFEST_V3": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v3.json", + ArtifactsTypes.MANIFEST_V3, DestinationTables.MANIFEST_V3, ManifestV3), + "RUN_RESULTS_V3": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v3.json", + ArtifactsTypes.RUN_RESULTS_V3, DestinationTables.RUN_RESULTS_V3, RunResultsV3), + "SOURCES_V3": ArtifactInfo("https://schemas.getdbt.com/dbt/sources/v3.json", + ArtifactsTypes.SOURCES_V3, DestinationTables.SOURCES_V3, SourcesV3), + # V4 + "MANIFEST_V4": ArtifactInfo("https://schemas.getdbt.com/dbt/manifest/v4.json", + ArtifactsTypes.MANIFEST_V4, DestinationTables.MANIFEST_V4, ManifestV4), + "RUN_RESULTS_V4": ArtifactInfo("https://schemas.getdbt.com/dbt/run-results/v4.json", + ArtifactsTypes.RUN_RESULTS_V4, DestinationTables.RUN_RESULTS_V4, RunResultsV4), +} diff --git a/terraform/example/main.tf b/terraform/example/main.tf index a57a2c5..e5634df 100644 --- a/terraform/example/main.tf +++ b/terraform/example/main.tf @@ -6,7 +6,7 @@ module "dbt_artifacts_loader" { delete_on_destroy = true - docker_image = "gcr.io/${var.project_id}/dbt-artifacts-loader:v1.1.0-rc2" + docker_image = "gcr.io/${var.project_id}/dbt-artifacts-loader:v1.2.0-dev2" labels = { app = "dbt-artifacts-loader" diff --git a/tests/dbt/test_model_factory.py b/tests/dbt/test_model_factory.py deleted file mode 100644 index 1a9d4ce..0000000 --- a/tests/dbt/test_model_factory.py +++ /dev/null @@ -1,45 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# -import unittest - -from dbt_artifacts_loader.dbt.utils import ArtifactsTypes -from dbt_artifacts_loader.dbt.v1.catalog import CatalogV1 -from dbt_artifacts_loader.dbt.v1.manifest import ManifestV1 -from dbt_artifacts_loader.dbt.v1.run_results import RunResultsV1 -from dbt_artifacts_loader.dbt.v1.sources import SourcesV1 -from dbt_artifacts_loader.dbt.v2.manifest import ManifestV2 -from dbt_artifacts_loader.dbt.v2.run_results import RunResultsV2 -from dbt_artifacts_loader.dbt.model_factory import get_model_class - - -class TestModelFactory(unittest.TestCase): - - def test_get_model_class(self): - test_sets = [ - # v1 - (ArtifactsTypes.CATALOG_V1, CatalogV1), - (ArtifactsTypes.MANIFEST_V1, ManifestV1), - (ArtifactsTypes.RUN_RESULTS_V1, RunResultsV1), - (ArtifactsTypes.SOURCES_V1, SourcesV1), - # v2 - (ArtifactsTypes.MANIFEST_V2, ManifestV2), - (ArtifactsTypes.RUN_RESULTS_V2, RunResultsV2), - ] - for (artifact_type, expected_class) in test_sets: - klass = get_model_class(artifact_type=artifact_type) - self.assertEqual(klass, expected_class) diff --git a/tests/dbt/test_utils.py b/tests/dbt/test_utils.py index 35c086b..a7390e3 100644 --- a/tests/dbt/test_utils.py +++ b/tests/dbt/test_utils.py @@ -19,9 +19,21 @@ import unittest import json - +from dbt_artifacts_loader.dbt.v1.catalog import CatalogV1 +from dbt_artifacts_loader.dbt.v1.manifest import ManifestV1 +from dbt_artifacts_loader.dbt.v1.run_results import RunResultsV1 +from dbt_artifacts_loader.dbt.v1.sources import SourcesV1 +from dbt_artifacts_loader.dbt.v2.manifest import ManifestV2 +from dbt_artifacts_loader.dbt.v2.run_results import RunResultsV2 +from dbt_artifacts_loader.dbt.v2.sources import SourcesV2 +from dbt_artifacts_loader.dbt.v3.manifest import ManifestV3 +from dbt_artifacts_loader.dbt.v3.run_results import RunResultsV3 +from dbt_artifacts_loader.dbt.v3.sources import SourcesV3 +from dbt_artifacts_loader.dbt.v4.manifest import ManifestV4 +from dbt_artifacts_loader.dbt.v4.run_results import RunResultsV4 +from dbt_artifacts_loader.dbt.version_map import ArtifactsTypes from dbt_artifacts_loader.utils import get_project_root -from dbt_artifacts_loader.dbt.utils import get_dbt_schema_version +from dbt_artifacts_loader.dbt.utils import get_dbt_schema_version, get_model_class class TestDbtUtils(unittest.TestCase): @@ -61,3 +73,26 @@ def test_get_dbt_schema_version(self): artifact_json = json.load(fp) dbt_schema_version = get_dbt_schema_version(artifact_json=artifact_json) self.assertEqual(dbt_schema_version, expected_dbt_schema_version) + + def test_get_model_class(self): + test_sets = [ + # v1 + (ArtifactsTypes.CATALOG_V1, CatalogV1), + (ArtifactsTypes.MANIFEST_V1, ManifestV1), + (ArtifactsTypes.RUN_RESULTS_V1, RunResultsV1), + (ArtifactsTypes.SOURCES_V1, SourcesV1), + # v2 + (ArtifactsTypes.MANIFEST_V2, ManifestV2), + (ArtifactsTypes.RUN_RESULTS_V2, RunResultsV2), + (ArtifactsTypes.SOURCES_V2, SourcesV2), + # v3 + (ArtifactsTypes.MANIFEST_V3, ManifestV3), + (ArtifactsTypes.RUN_RESULTS_V3, RunResultsV3), + (ArtifactsTypes.SOURCES_V3, SourcesV3), + # v4 + (ArtifactsTypes.MANIFEST_V4, ManifestV4), + (ArtifactsTypes.RUN_RESULTS_V4, RunResultsV4), + ] + for (artifact_type, expected_class) in test_sets: + klass = get_model_class(artifact_type=artifact_type) + self.assertEqual(klass, expected_class)