diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 01518f4b..12796f39 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ repos: - id: ruff args: [--fix, --exit-non-zero-on-fix] - id: ruff-format - rev: v0.11.12 + rev: v0.15.6 - repo: local hooks: - id: mypy @@ -14,7 +14,7 @@ repos: entry: uv run mypy --non-interactive --install-types cognite pass_filenames: false - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: check-ast - id: debug-statements diff --git a/cognite/extractorutils/base.py b/cognite/extractorutils/base.py index fff11a76..20f1751d 100644 --- a/cognite/extractorutils/base.py +++ b/cognite/extractorutils/base.py @@ -28,11 +28,11 @@ from types import TracebackType from typing import Any, Generic, Literal, TypeVar -from dotenv import load_dotenv - from cognite.client import CogniteClient from cognite.client.data_classes import ExtractionPipeline, ExtractionPipelineRun from cognite.client.exceptions import CogniteAPIError +from dotenv import load_dotenv + from cognite.extractorutils.configtools import ( BaseConfig, ConfigResolver, diff --git a/cognite/extractorutils/configtools/elements.py b/cognite/extractorutils/configtools/elements.py index 3a1f18fc..aa077818 100644 --- a/cognite/extractorutils/configtools/elements.py +++ b/cognite/extractorutils/configtools/elements.py @@ -28,9 +28,6 @@ from urllib.parse import urljoin, urlparse import yaml -from prometheus_client import REGISTRY, start_http_server -from typing_extensions import Self - from cognite.client import ClientConfig, CogniteClient from cognite.client.credentials import ( CredentialProvider, @@ -38,6 +35,9 @@ OAuthClientCredentials, ) from cognite.client.data_classes import Asset, DataSet, ExtractionPipeline +from prometheus_client import REGISTRY, start_http_server +from typing_extensions import Self + from cognite.extractorutils._inner_util import resolve_log_level_for_httpx from cognite.extractorutils.configtools._util import _load_certificate_data from cognite.extractorutils.exceptions import InvalidConfigError diff --git a/cognite/extractorutils/configtools/loaders.py b/cognite/extractorutils/configtools/loaders.py index ae7bb332..d57ea227 100644 --- a/cognite/extractorutils/configtools/loaders.py +++ b/cognite/extractorutils/configtools/loaders.py @@ -34,9 +34,9 @@ from azure.core.exceptions import HttpResponseError, ResourceNotFoundError, ServiceRequestError from azure.identity import ClientSecretCredential, DefaultAzureCredential from azure.keyvault.secrets import SecretClient +from cognite.client import CogniteClient from yaml.scanner import ScannerError -from cognite.client import CogniteClient from cognite.extractorutils.configtools._util import _to_snake_case from cognite.extractorutils.configtools.elements import ( BaseConfig, diff --git a/cognite/extractorutils/metrics.py b/cognite/extractorutils/metrics.py index ff91581b..124fbaa6 100644 --- a/cognite/extractorutils/metrics.py +++ b/cognite/extractorutils/metrics.py @@ -48,13 +48,13 @@ def __init__(self): import arrow import psutil +from cognite.client import CogniteClient +from cognite.client.data_classes import Asset, TimeSeries +from cognite.client.exceptions import CogniteDuplicatedError from prometheus_client import Gauge, Info, Metric from prometheus_client.core import REGISTRY from prometheus_client.exposition import basic_auth_handler, delete_from_gateway, pushadd_to_gateway -from cognite.client import CogniteClient -from cognite.client.data_classes import Asset, TimeSeries -from cognite.client.exceptions import CogniteDuplicatedError from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader.time_series import DataPointList, TimeSeriesUploadQueue from cognite.extractorutils.util import EitherId diff --git a/cognite/extractorutils/statestore/hashing.py b/cognite/extractorutils/statestore/hashing.py index fd4b945e..053fc472 100644 --- a/cognite/extractorutils/statestore/hashing.py +++ b/cognite/extractorutils/statestore/hashing.py @@ -15,10 +15,10 @@ from typing import Any import orjson - from cognite.client import CogniteClient from cognite.client.data_classes import Row from cognite.client.exceptions import CogniteAPIError + from cognite.extractorutils._inner_util import _DecimalDecoder, _DecimalEncoder from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.util import cognite_exceptions, retry diff --git a/cognite/extractorutils/statestore/watermark.py b/cognite/extractorutils/statestore/watermark.py index 06cd5802..9f8269df 100644 --- a/cognite/extractorutils/statestore/watermark.py +++ b/cognite/extractorutils/statestore/watermark.py @@ -41,6 +41,7 @@ from cognite.client import CogniteClient from cognite.client.exceptions import CogniteAPIError + from cognite.extractorutils._inner_util import _DecimalDecoder, _DecimalEncoder from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader import DataPointList diff --git a/cognite/extractorutils/unstable/configuration/loaders.py b/cognite/extractorutils/unstable/configuration/loaders.py index b09bb0af..b718fd87 100644 --- a/cognite/extractorutils/unstable/configuration/loaders.py +++ b/cognite/extractorutils/unstable/configuration/loaders.py @@ -8,10 +8,10 @@ from pathlib import Path from typing import TextIO, TypeVar -from pydantic import ValidationError - from cognite.client import CogniteClient from cognite.client.exceptions import CogniteAPIError +from pydantic import ValidationError + from cognite.extractorutils.configtools.loaders import _load_yaml_dict_raw from cognite.extractorutils.exceptions import InvalidConfigError as OldInvalidConfigError from cognite.extractorutils.unstable.configuration.exceptions import InvalidConfigError diff --git a/cognite/extractorutils/unstable/configuration/models.py b/cognite/extractorutils/unstable/configuration/models.py index 5729592d..60f919f8 100644 --- a/cognite/extractorutils/unstable/configuration/models.py +++ b/cognite/extractorutils/unstable/configuration/models.py @@ -11,12 +11,6 @@ from time import sleep from typing import Annotated, Any, Literal, TypeVar -from humps import kebabize -from prometheus_client import REGISTRY, start_http_server -from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler -from pydantic_core import CoreSchema, core_schema -from typing_extensions import assert_never - from cognite.client import CogniteClient from cognite.client.config import ClientConfig from cognite.client.credentials import ( @@ -25,6 +19,12 @@ OAuthClientCredentials, ) from cognite.client.data_classes import Asset, DataSet +from humps import kebabize +from prometheus_client import REGISTRY, start_http_server +from pydantic import BaseModel, ConfigDict, Field, GetCoreSchemaHandler +from pydantic_core import CoreSchema, core_schema +from typing_extensions import assert_never + from cognite.extractorutils.configtools._util import _load_certificate_data from cognite.extractorutils.exceptions import InvalidConfigError from cognite.extractorutils.metrics import AbstractMetricsPusher, CognitePusher, PrometheusPusher diff --git a/cognite/extractorutils/unstable/core/checkin_worker.py b/cognite/extractorutils/unstable/core/checkin_worker.py index a65a9c31..568d3c57 100644 --- a/cognite/extractorutils/unstable/core/checkin_worker.py +++ b/cognite/extractorutils/unstable/core/checkin_worker.py @@ -16,10 +16,10 @@ from threading import RLock from time import sleep -from requests import Response - from cognite.client import CogniteClient from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteConnectionError +from requests import Response + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.unstable.configuration.models import ConfigRevision from cognite.extractorutils.unstable.core._dto import ( diff --git a/cognite/extractorutils/unstable/core/runtime.py b/cognite/extractorutils/unstable/core/runtime.py index da0a2068..4cedd4af 100644 --- a/cognite/extractorutils/unstable/core/runtime.py +++ b/cognite/extractorutils/unstable/core/runtime.py @@ -38,15 +38,15 @@ def main() -> None: from typing import Any, Generic, TypeVar from uuid import uuid4 -from requests.exceptions import ConnectionError as RequestsConnectionError -from typing_extensions import assert_never - from cognite.client import CogniteClient from cognite.client.exceptions import ( CogniteAPIError, CogniteAuthError, CogniteConnectionError, ) +from requests.exceptions import ConnectionError as RequestsConnectionError +from typing_extensions import assert_never + from cognite.extractorutils.metrics import BaseMetrics, MetricsType from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.unstable.configuration.exceptions import InvalidArgumentError, InvalidConfigError diff --git a/cognite/extractorutils/uploader/_base.py b/cognite/extractorutils/uploader/_base.py index a4cb3ad7..2d4b785f 100644 --- a/cognite/extractorutils/uploader/_base.py +++ b/cognite/extractorutils/uploader/_base.py @@ -20,8 +20,8 @@ from typing import Any from arrow import Arrow - from cognite.client import CogniteClient + from cognite.extractorutils._inner_util import _resolve_log_level from cognite.extractorutils.threading import CancellationToken diff --git a/cognite/extractorutils/uploader/assets.py b/cognite/extractorutils/uploader/assets.py index 5da31fa3..e2949ccc 100644 --- a/cognite/extractorutils/uploader/assets.py +++ b/cognite/extractorutils/uploader/assets.py @@ -23,6 +23,7 @@ from cognite.client import CogniteClient from cognite.client.data_classes.assets import Asset from cognite.client.exceptions import CogniteDuplicatedError + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader._base import ( RETRIES, diff --git a/cognite/extractorutils/uploader/data_modeling.py b/cognite/extractorutils/uploader/data_modeling.py index 113f2f33..0863166f 100644 --- a/cognite/extractorutils/uploader/data_modeling.py +++ b/cognite/extractorutils/uploader/data_modeling.py @@ -8,6 +8,7 @@ from cognite.client import CogniteClient from cognite.client.data_classes.data_modeling import EdgeApply, NodeApply + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader._base import ( RETRIES, diff --git a/cognite/extractorutils/uploader/events.py b/cognite/extractorutils/uploader/events.py index 578c117f..5cba1d7d 100644 --- a/cognite/extractorutils/uploader/events.py +++ b/cognite/extractorutils/uploader/events.py @@ -22,6 +22,7 @@ from cognite.client import CogniteClient from cognite.client.data_classes import Event from cognite.client.exceptions import CogniteDuplicatedError + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader._base import ( RETRIES, diff --git a/cognite/extractorutils/uploader/files.py b/cognite/extractorutils/uploader/files.py index be4f6ac6..6a7c83c7 100644 --- a/cognite/extractorutils/uploader/files.py +++ b/cognite/extractorutils/uploader/files.py @@ -29,10 +29,6 @@ ) from urllib.parse import ParseResult, urlparse -from h11._util import LocalProtocolError -from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream -from requests.utils import super_len - from cognite.client import CogniteClient from cognite.client.data_classes import FileMetadata, FileMetadataUpdate from cognite.client.data_classes.data_modeling import NodeId @@ -40,6 +36,10 @@ CogniteExtractorFileApply, ) from cognite.client.utils._identifier import IdentifierSequence +from h11._util import LocalProtocolError +from httpx import URL, Client, Headers, Request, StreamConsumed, SyncByteStream +from requests.utils import super_len + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader._base import ( RETRIES, @@ -383,7 +383,7 @@ def _upload_empty_file( self, file_meta: FileMetadataOrCogniteExtractorFile, ) -> None: - file_meta_response, url = self._upload_only_metadata(file_meta) + _file_meta_response, url = self._upload_only_metadata(file_meta) self._upload_only_file_reference(file_meta, url) diff --git a/cognite/extractorutils/uploader/raw.py b/cognite/extractorutils/uploader/raw.py index 8161d1f9..1e3e709e 100644 --- a/cognite/extractorutils/uploader/raw.py +++ b/cognite/extractorutils/uploader/raw.py @@ -22,9 +22,9 @@ import arrow from arrow import Arrow - from cognite.client import CogniteClient from cognite.client.data_classes import Row + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader._base import ( RETRIES, diff --git a/cognite/extractorutils/uploader/time_series.py b/cognite/extractorutils/uploader/time_series.py index ce4e8b30..019b60cd 100644 --- a/cognite/extractorutils/uploader/time_series.py +++ b/cognite/extractorutils/uploader/time_series.py @@ -33,6 +33,7 @@ from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import CogniteExtractorTimeSeriesApply from cognite.client.data_classes.data_modeling.instances import DirectRelationReference from cognite.client.exceptions import CogniteDuplicatedError, CogniteNotFoundError + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.uploader._base import ( RETRIES, diff --git a/cognite/extractorutils/uploader_extractor.py b/cognite/extractorutils/uploader_extractor.py index 98e91c6d..4f7d5de1 100644 --- a/cognite/extractorutils/uploader_extractor.py +++ b/cognite/extractorutils/uploader_extractor.py @@ -23,9 +23,9 @@ from types import TracebackType from typing import Any, TypeVar +from cognite.client import CogniteClient from more_itertools import peekable -from cognite.client import CogniteClient from cognite.extractorutils.base import Extractor from cognite.extractorutils.configtools import BaseConfig, TimeIntervalConfig from cognite.extractorutils.metrics import BaseMetrics @@ -176,7 +176,7 @@ def __enter__(self) -> "UploaderExtractor": """ super().__enter__() - queue_config = self.config.queues if self.config.queues else QueueConfigClass() + queue_config = self.config.queues or QueueConfigClass() self.event_queue = EventUploadQueue( self.cognite_client, max_queue_size=queue_config.event_size, diff --git a/cognite/extractorutils/uploader_types.py b/cognite/extractorutils/uploader_types.py index 9bdbd1ce..2ab7c5af 100644 --- a/cognite/extractorutils/uploader_types.py +++ b/cognite/extractorutils/uploader_types.py @@ -10,6 +10,7 @@ from cognite.client.data_classes import Event as _Event from cognite.client.data_classes import Row as _Row from cognite.client.data_classes.data_modeling import NodeId + from cognite.extractorutils.uploader.time_series import DataPoint diff --git a/cognite/extractorutils/util.py b/cognite/extractorutils/util.py index e4ad8745..ffe64e0d 100644 --- a/cognite/extractorutils/util.py +++ b/cognite/extractorutils/util.py @@ -19,7 +19,7 @@ import io import logging import random -from collections.abc import Callable, Generator, Iterable +from collections.abc import Callable, Generator, Iterable, Sequence from datetime import datetime, timezone from functools import partial, wraps from io import RawIOBase @@ -27,13 +27,13 @@ from time import time from typing import Any, TypeVar -from decorator import decorator - from cognite.client import CogniteClient from cognite.client._api.assets import AssetsAPI from cognite.client._api.time_series import TimeSeriesAPI from cognite.client.data_classes import Asset, ExtractionPipelineRun, TimeSeries from cognite.client.exceptions import CogniteAPIError, CogniteException, CogniteFileUploadError, CogniteNotFoundError +from decorator import decorator + from cognite.extractorutils.threading import CancellationToken @@ -400,8 +400,8 @@ def retry( @decorator def retry_decorator(f: Callable[..., _T2], *fargs: Any, **fkwargs: Any) -> _T2: # noqa: ANN401 - args = fargs if fargs else [] - kwargs = fkwargs if fkwargs else {} + args: Sequence[Any] = fargs or [] + kwargs: dict[str, Any] = fkwargs or {} return _retry_internal( partial(f, *args, **kwargs), diff --git a/pyproject.toml b/pyproject.toml index 945a8e82..29d43f27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,16 +36,10 @@ dependencies = [ "simple-winservice>=0.1.0 ; sys_platform == 'win32'", ] -[project.urls] -repository = "https://github.com/cognitedata/python-extractor-utils" - -[project.optional-dependencies] -experimental = ["cognite-sdk-experimental"] - -[tool.uv] -dev-dependencies = [ - "mypy>=1.14.1", - "ruff>=0.9.0", +[dependency-groups] +dev = [ + "mypy>=1.19.1", + "ruff>=0.15.0", "pytest>=8.0.0", "pytest-cov>=6.0.0", "sphinx>=7.0.0", @@ -65,6 +59,12 @@ dev-dependencies = [ "requests-mock>=1.12.1", ] +[project.urls] +repository = "https://github.com/cognitedata/python-extractor-utils" + +[project.optional-dependencies] +experimental = ["cognite-sdk-experimental"] + [tool.ruff] exclude = [ ".git", @@ -123,8 +123,5 @@ only-include = ["cognite"] requires = ["hatchling"] build-backend = "hatchling.build" -[dependency-groups] -dev = [] - [project.scripts] simple-extractor = "cognite.examples.unstable.extractors.simple_extractor.main:main" diff --git a/tests/conftest.py b/tests/conftest.py index a20f2127..be7cf1eb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,13 +5,13 @@ from enum import Enum import pytest -from prometheus_client.core import REGISTRY - from cognite.client import CogniteClient from cognite.client.config import ClientConfig from cognite.client.credentials import OAuthClientCredentials from cognite.client.data_classes.data_modeling import NodeId from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError +from prometheus_client.core import REGISTRY + from cognite.extractorutils import metrics NUM_NODES = 5000 diff --git a/tests/test_unstable/conftest.py b/tests/test_unstable/conftest.py index d143f0ff..91edce00 100644 --- a/tests/test_unstable/conftest.py +++ b/tests/test_unstable/conftest.py @@ -9,11 +9,11 @@ import pytest import requests_mock -from prometheus_client.core import Counter - from cognite.client import CogniteClient from cognite.client.config import ClientConfig from cognite.client.credentials import OAuthClientCredentials +from prometheus_client.core import Counter + from cognite.extractorutils.metrics import BaseMetrics from cognite.extractorutils.unstable.configuration.models import ( ConnectionConfig, diff --git a/tests/test_unstable/test_base.py b/tests/test_unstable/test_base.py index 4721c79b..4aa81ac0 100644 --- a/tests/test_unstable/test_base.py +++ b/tests/test_unstable/test_base.py @@ -7,9 +7,9 @@ from pathlib import Path import pytest - from cognite.client import CogniteClient from cognite.client.exceptions import CogniteNotFoundError + from cognite.extractorutils.metrics import CognitePusher, PrometheusPusher from cognite.extractorutils.statestore.watermark import LocalStateStore, RawStateStore from cognite.extractorutils.unstable.configuration.loaders import ConfigFormat, load_io @@ -119,7 +119,7 @@ def test_get_current_statestore_raises_before_start() -> None: Tests that calling get_current_statestore before the extractor's __enter__ method is called raises a ValueError. """ - with pytest.raises(ValueError, match="No state store singleton created. Have a state store been loaded?"): + with pytest.raises(ValueError, match=r"No state store singleton created. Have a state store been loaded\?"): TestExtractor.get_current_statestore() diff --git a/tests/test_unstable/test_configuration.py b/tests/test_unstable/test_configuration.py index 43a7588a..be1393b4 100644 --- a/tests/test_unstable/test_configuration.py +++ b/tests/test_unstable/test_configuration.py @@ -3,10 +3,10 @@ from unittest.mock import Mock import pytest -from pydantic import Field - from cognite.client.credentials import OAuthClientCredentials from cognite.client.data_classes import DataSet +from pydantic import Field + from cognite.extractorutils.exceptions import InvalidConfigError from cognite.extractorutils.unstable.configuration.loaders import ConfigFormat, load_io from cognite.extractorutils.unstable.configuration.models import ( diff --git a/tests/test_unstable/test_is_dev_cluster.py b/tests/test_unstable/test_is_dev_cluster.py index 102eb366..46b81616 100644 --- a/tests/test_unstable/test_is_dev_cluster.py +++ b/tests/test_unstable/test_is_dev_cluster.py @@ -1,7 +1,6 @@ import os import pytest - from cognite.client import CogniteClient diff --git a/tests/tests_integration/test_assets_integration.py b/tests/tests_integration/test_assets_integration.py index 099ac214..96ccc614 100644 --- a/tests/tests_integration/test_assets_integration.py +++ b/tests/tests_integration/test_assets_integration.py @@ -16,10 +16,10 @@ import random import pytest -from conftest import ETestType, ParamTest - from cognite.client import CogniteClient from cognite.client.data_classes.assets import Asset +from conftest import ETestType, ParamTest + from cognite.extractorutils.uploader.assets import AssetUploadQueue diff --git a/tests/tests_integration/test_cdm_time_series_integration.py b/tests/tests_integration/test_cdm_time_series_integration.py index 98df0d3b..0b66fbf8 100644 --- a/tests/tests_integration/test_cdm_time_series_integration.py +++ b/tests/tests_integration/test_cdm_time_series_integration.py @@ -8,13 +8,13 @@ import pytest from _pytest.logging import LogCaptureFixture -from conftest import ETestType, ParamTest - from cognite.client import CogniteClient from cognite.client.data_classes import StatusCode from cognite.client.data_classes.data_modeling import NodeApply, NodeId from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import CogniteExtractorTimeSeriesApply from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError +from conftest import ETestType, ParamTest + from cognite.extractorutils.uploader.time_series import ( MAX_DATAPOINT_STRING_BYTES, MAX_DATAPOINT_VALUE, diff --git a/tests/tests_integration/test_configtools_integration.py b/tests/tests_integration/test_configtools_integration.py index b1e1ed2f..2308e841 100644 --- a/tests/tests_integration/test_configtools_integration.py +++ b/tests/tests_integration/test_configtools_integration.py @@ -19,6 +19,7 @@ from cognite.client import CogniteClient from cognite.client.data_classes import DataSet, ExtractionPipelineConfigWrite, ExtractionPipelineWrite from cognite.client.exceptions import CogniteDuplicatedError + from cognite.extractorutils.configtools import BaseConfig, load_yaml from cognite.extractorutils.configtools.loaders import ConfigResolver diff --git a/tests/tests_integration/test_data_modeling_queue.py b/tests/tests_integration/test_data_modeling_queue.py index affed12f..3184f796 100644 --- a/tests/tests_integration/test_data_modeling_queue.py +++ b/tests/tests_integration/test_data_modeling_queue.py @@ -2,9 +2,6 @@ from time import sleep import pytest -from conftest import NUM_EDGES, NUM_NODES, ETestType, ParamTest -from faker import Faker - from cognite.client import CogniteClient from cognite.client.data_classes.data_modeling import ( EdgeApply, @@ -12,6 +9,9 @@ NodeOrEdgeData, ViewId, ) +from conftest import NUM_EDGES, NUM_NODES, ETestType, ParamTest +from faker import Faker + from cognite.extractorutils.uploader.data_modeling import InstanceUploadQueue fake = Faker() diff --git a/tests/tests_integration/test_events_integration.py b/tests/tests_integration/test_events_integration.py index caa3eee5..6bdb85e9 100644 --- a/tests/tests_integration/test_events_integration.py +++ b/tests/tests_integration/test_events_integration.py @@ -16,10 +16,10 @@ import random import pytest -from conftest import ETestType, ParamTest - from cognite.client import CogniteClient from cognite.client.data_classes import Event +from conftest import ETestType, ParamTest + from cognite.extractorutils.uploader.events import EventUploadQueue diff --git a/tests/tests_integration/test_file_integration.py b/tests/tests_integration/test_file_integration.py index 65dda9c6..a54d0462 100644 --- a/tests/tests_integration/test_file_integration.py +++ b/tests/tests_integration/test_file_integration.py @@ -22,8 +22,6 @@ import jsonlines import pytest -from conftest import ETestType, ParamTest - from cognite.client import CogniteClient from cognite.client.data_classes import FileMetadata from cognite.client.data_classes.data_modeling import NodeId @@ -31,6 +29,8 @@ CogniteExtractorFile, CogniteExtractorFileApply, ) +from conftest import ETestType, ParamTest + from cognite.extractorutils.uploader.files import ( BytesUploadQueue, FileUploadQueue, diff --git a/tests/tests_integration/test_metrics_integration.py b/tests/tests_integration/test_metrics_integration.py index 3f97967a..b2fa53f3 100644 --- a/tests/tests_integration/test_metrics_integration.py +++ b/tests/tests_integration/test_metrics_integration.py @@ -12,11 +12,11 @@ from typing import Any import pytest +from cognite.client import CogniteClient +from cognite.client.exceptions import CogniteNotFoundError from prometheus_client import Counter, Gauge from prometheus_client.core import REGISTRY -from cognite.client import CogniteClient -from cognite.client.exceptions import CogniteNotFoundError from cognite.extractorutils.metrics import CognitePusher logger = logging.getLogger(__name__) diff --git a/tests/tests_integration/test_raw_integration.py b/tests/tests_integration/test_raw_integration.py index f48e6b72..a38e291c 100644 --- a/tests/tests_integration/test_raw_integration.py +++ b/tests/tests_integration/test_raw_integration.py @@ -17,10 +17,10 @@ import time import pytest -from conftest import ETestType, ParamTest - from cognite.client import CogniteClient from cognite.client.data_classes import Row +from conftest import ETestType, ParamTest + from cognite.extractorutils.uploader import RawUploadQueue diff --git a/tests/tests_integration/test_time_series_integration.py b/tests/tests_integration/test_time_series_integration.py index 06a4a458..b525a4a9 100644 --- a/tests/tests_integration/test_time_series_integration.py +++ b/tests/tests_integration/test_time_series_integration.py @@ -19,10 +19,10 @@ from typing import Any import pytest -from conftest import ETestType, ParamTest - from cognite.client import CogniteClient from cognite.client.data_classes import StatusCode, TimeSeries +from conftest import ETestType, ParamTest + from cognite.extractorutils.uploader import TimeSeriesUploadQueue diff --git a/tests/tests_unit/test_base.py b/tests/tests_unit/test_base.py index 2312f93d..2dc7d9e2 100644 --- a/tests/tests_unit/test_base.py +++ b/tests/tests_unit/test_base.py @@ -18,9 +18,9 @@ from unittest.mock import Mock, patch import pytest - from cognite.client import CogniteClient from cognite.client.data_classes import ExtractionPipeline, ExtractionPipelineRun + from cognite.extractorutils import Extractor from cognite.extractorutils.configtools import BaseConfig, StateStoreConfig from cognite.extractorutils.statestore import LocalStateStore, NoStateStore diff --git a/tests/tests_unit/test_cdf_upload_queues.py b/tests/tests_unit/test_cdf_upload_queues.py index 7c53107c..383514be 100644 --- a/tests/tests_unit/test_cdf_upload_queues.py +++ b/tests/tests_unit/test_cdf_upload_queues.py @@ -19,13 +19,13 @@ from typing import Any, BinaryIO from unittest.mock import Mock, patch -from httpx import URL, Request - from cognite.client.data_classes import Event, FileMetadata, Row from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import ( CogniteExtractorFileApply, ) from cognite.client.testing import CogniteClientMock +from httpx import URL, Request + from cognite.extractorutils.uploader import ( EventUploadQueue, IOFileUploadQueue, diff --git a/tests/tests_unit/test_configtools.py b/tests/tests_unit/test_configtools.py index 3bd2676d..60e813ce 100644 --- a/tests/tests_unit/test_configtools.py +++ b/tests/tests_unit/test_configtools.py @@ -26,10 +26,10 @@ import pytest import yaml -from faker import Faker - from cognite.client import CogniteClient from cognite.client.credentials import OAuthClientCredentials +from faker import Faker + from cognite.extractorutils.configtools import ( BaseConfig, CogniteConfig, diff --git a/tests/tests_unit/test_extraction_pipelines.py b/tests/tests_unit/test_extraction_pipelines.py index 0726c548..41ab1636 100644 --- a/tests/tests_unit/test_extraction_pipelines.py +++ b/tests/tests_unit/test_extraction_pipelines.py @@ -15,8 +15,8 @@ import time import pytest - from cognite.client.testing import monkeypatch_cognite_client + from cognite.extractorutils.util import add_extraction_pipeline with monkeypatch_cognite_client() as m_client: diff --git a/tests/tests_unit/test_metrics.py b/tests/tests_unit/test_metrics.py index 6d2f0505..89c44540 100644 --- a/tests/tests_unit/test_metrics.py +++ b/tests/tests_unit/test_metrics.py @@ -18,11 +18,11 @@ import arrow import pytest -from prometheus_client import Gauge - from cognite.client import CogniteClient from cognite.client.data_classes import Asset from cognite.client.exceptions import CogniteDuplicatedError, CogniteNotFoundError +from prometheus_client import Gauge + from cognite.extractorutils.metrics import CognitePusher, safe_get diff --git a/tests/tests_unit/test_statestore.py b/tests/tests_unit/test_statestore.py index bb20b9c3..40f466c2 100644 --- a/tests/tests_unit/test_statestore.py +++ b/tests/tests_unit/test_statestore.py @@ -22,10 +22,10 @@ from unittest.mock import Mock, patch import pytest - from cognite.client import CogniteClient from cognite.client.data_classes import Row from cognite.client.exceptions import CogniteAPIError + from cognite.extractorutils.statestore import LocalHashStateStore, LocalStateStore, NoStateStore, RawStateStore from cognite.extractorutils.uploader import TimeSeriesUploadQueue diff --git a/tests/tests_unit/test_uploader_extractor.py b/tests/tests_unit/test_uploader_extractor.py index 3b1e45ac..c4660085 100644 --- a/tests/tests_unit/test_uploader_extractor.py +++ b/tests/tests_unit/test_uploader_extractor.py @@ -18,6 +18,7 @@ from cognite.client.data_classes import Row from cognite.client.data_classes.data_modeling import NodeApply, NodeId from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import CogniteExtractorTimeSeriesApply + from cognite.extractorutils.uploader import ( CDMTimeSeriesUploadQueue, EventUploadQueue, diff --git a/tests/tests_unit/test_util.py b/tests/tests_unit/test_util.py index cc7d876f..5b74a512 100644 --- a/tests/tests_unit/test_util.py +++ b/tests/tests_unit/test_util.py @@ -19,9 +19,9 @@ import httpx import pytest import requests - from cognite.client.data_classes import Asset, TimeSeries from cognite.client.exceptions import CogniteAPIError, CogniteFileUploadError, CogniteNotFoundError + from cognite.extractorutils.threading import CancellationToken from cognite.extractorutils.util import ( EitherId,