Add unit tests for previously-untested modules; fix DynamicConfig.get_config()

CHANGELOG.md

@@ -1,4 +1,5 @@
 * Support the timezone-carrying types `TzDate`, `TzDatetime` and `TzTimestamp` in query results and parameters (reading them previously raised `AttributeError`)
+* Fixed `DynamicConfigClient.get_config()` raising `TypeError` when parsing the server response
 * Drop support for Python 3.8 and 3.9; the minimum supported version is now Python 3.10
 * Accept native `datetime.datetime` values for `Datetime` and `Datetime64` query parameters (previously only integer seconds since the epoch were accepted)
 

tests/unit/test_dynamic_config.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from ydb import issues, operation
+from ydb._apis import ydb_operation
+from ydb.draft import dynamic_config as dc
+from ydb.draft import _apis
+
+
+class _FakeDriver:
+    def __init__(self, result="RESULT"):
+        self._result = result
+        self.calls = []
+
+    def __call__(self, *args):
+        self.calls.append(args)
+        return self._result
+
+    def future(self, *args):
+        self.calls.append(args)
+        return self._result
+
+
+def _get_config_response(version=7, cluster="cluster-1", config="cfg-text", status=issues.StatusCode.SUCCESS):
+    result = _apis.ydb_dynamic_config.GetConfigResult()
+    result.identity.version = version
+    result.identity.cluster = cluster
+    result.config = config
+    op = ydb_operation.Operation(status=status)
+    op.result.Pack(result)
+    return _apis.ydb_dynamic_config.GetConfigResponse(operation=op)
+
+
+def _get_node_labels_response(labels, status=issues.StatusCode.SUCCESS):
+    result = _apis.ydb_dynamic_config.GetNodeLabelsResult()
+    for label, value in labels.items():
+        result.labels.add(label=label, value=value)
+    op = ydb_operation.Operation(status=status)
+    op.result.Pack(result)
+    return _apis.ydb_dynamic_config.GetNodeLabelsResponse(operation=op)
+
+
+def test_dynamic_config_dataclass():
+    cfg = dc.DynamicConfig(3, "cl", "text", "extra")
+    assert cfg.version == 3
+    assert cfg.cluster == "cl"
+    assert cfg.config == "text"
+
+
+def test_node_labels_dataclass():
+    nl = dc.NodeLabels({"dc": "vla"})
+    assert nl.labels == {"dc": "vla"}
+
+
+def test_replace_and_set_config_request_factory():
+    replace = dc._replace_config_request_factory("cfg", True, False)
+    assert replace.config == "cfg"
+    assert replace.dry_run is True
+    assert replace.allow_unknown_fields is False
+
+    set_req = dc._set_config_request_factory("cfg2", False, True)
+    assert set_req.config == "cfg2"
+    assert set_req.dry_run is False
+    assert set_req.allow_unknown_fields is True
+
+
+def test_get_config_and_node_labels_request_factory():
+    assert dc._get_config_request_factory() is not None
+    node_req = dc._get_node_labels_request_factory(42)
+    assert node_req.node_id == 42
+
+
+def test_wrap_dynamic_config():
+    result = _apis.ydb_dynamic_config.GetConfigResult()
+    result.identity.version = 11
+    result.identity.cluster = "prod"
+    result.config = "yaml-config"
+    wrapped = dc._wrap_dynamic_config(result)
+    assert wrapped.version == 11
+    assert wrapped.cluster == "prod"
+    assert wrapped.config == "yaml-config"
+
+
+def test_wrap_get_config_response_success_and_error():
+    wrapped = dc._wrap_get_config_response(None, _get_config_response(version=9, cluster="c", config="cfg"))
+    assert isinstance(wrapped, dc.DynamicConfig)
+    assert wrapped.version == 9
+    assert wrapped.cluster == "c"
+    assert wrapped.config == "cfg"
+
+    with pytest.raises(issues.BadRequest):
+        dc._wrap_get_config_response(None, _get_config_response(status=issues.StatusCode.BAD_REQUEST))
+
+
+def test_wrap_node_labels():
+    result = _apis.ydb_dynamic_config.GetNodeLabelsResult()
+    result.labels.add(label="dc", value="vla")
+    result.labels.add(label="rack", value="7")
+    wrapped = dc._wrap_node_labels(result)
+    assert wrapped.labels == {"dc": "vla", "rack": "7"}
+
+
+def test_wrap_get_node_labels_response_success_and_error():
+    wrapped = dc._wrap_get_node_labels_response(None, _get_node_labels_response({"dc": "sas"}))
+    assert isinstance(wrapped, dc.NodeLabels)
+    assert wrapped.labels == {"dc": "sas"}
+
+    with pytest.raises(issues.PreconditionFailed):
+        dc._wrap_get_node_labels_response(
+            None, _get_node_labels_response({}, status=issues.StatusCode.PRECONDITION_FAILED)
+        )
+
+
+def test_base_client_dispatch():
+    driver = _FakeDriver()
+    client = dc.BaseDynamicConfigClient(driver)
+
+    client.replace_config("cfg", True, False)
+    request, stub, method, wrapper, settings = driver.calls[-1]
+    assert request.config == "cfg"
+    assert stub is _apis.DynamicConfigService.Stub
+    assert method == _apis.DynamicConfigService.ReplaceConfig
+    assert wrapper is operation.Operation
+
+    client.set_config("cfg2", False, True)
+    assert driver.calls[-1][2] == _apis.DynamicConfigService.SetConfig
+
+    client.get_config()
+    request, _, method, wrapper, _ = driver.calls[-1]
+    assert method == _apis.DynamicConfigService.GetConfig
+    assert wrapper is dc._wrap_get_config_response
+
+    client.get_node_labels(5)
+    request, _, method, wrapper, _ = driver.calls[-1]
+    assert request.node_id == 5
+    assert method == _apis.DynamicConfigService.GetNodeLabels
+    assert wrapper is dc._wrap_get_node_labels_response
+
+
+def test_async_client_dispatch():
+    driver = _FakeDriver()
+    client = dc.DynamicConfigClient(driver)
+
+    client.async_replace_config("cfg", True, False)
+    assert driver.calls[-1][2] == _apis.DynamicConfigService.ReplaceConfig
+
+    client.async_set_config("cfg", False, True)
+    assert driver.calls[-1][2] == _apis.DynamicConfigService.SetConfig
+
+    client.async_get_config()
+    assert driver.calls[-1][3] is dc._wrap_get_config_response
+
+    client.async_get_node_labels(9)
+    request, _, method, wrapper, _ = driver.calls[-1]
+    assert request.node_id == 9
+    assert wrapper is dc._wrap_get_node_labels_response

tests/unit/test_export.py

@@ -0,0 +1,194 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from ydb import export, issues, _apis
+from ydb._grpc.common.protos import ydb_export_pb2
+from ydb._grpc.common import ydb_export_v1_pb2_grpc
+
+
+class _FakeDriver:
+    def __init__(self, result="RESULT"):
+        self._result = result
+        self.calls = []
+
+    def __call__(self, *args):
+        self.calls.append(args)
+        return self._result
+
+    def future(self, *args):
+        self.calls.append(args)
+        return self._result
+
+
+def _s3_response(progress=ydb_export_pb2.ExportProgress.PROGRESS_DONE, operation_id="op-1"):
+    op = _apis.ydb_operation.Operation(id=operation_id, status=issues.StatusCode.SUCCESS)
+    op.metadata.Pack(ydb_export_pb2.ExportToS3Metadata(progress=progress))
+    return _apis.ydb_operation.GetOperationResponse(operation=op)
+
+
+def _yt_response(progress=ydb_export_pb2.ExportProgress.PROGRESS_TRANSFER_DATA, operation_id="op-2"):
+    op = _apis.ydb_operation.Operation(id=operation_id, status=issues.StatusCode.SUCCESS)
+    op.metadata.Pack(ydb_export_pb2.ExportToYtMetadata(progress=progress))
+    return _apis.ydb_operation.GetOperationResponse(operation=op)
+
+
+def test_export_progress_enum_initialized():
+    # PROGRESS_* proto values are mapped onto the ExportProgress enum at import.
+    assert export._progresses[ydb_export_pb2.ExportProgress.PROGRESS_DONE] is export.ExportProgress.DONE
+    assert export._progresses[ydb_export_pb2.ExportProgress.PROGRESS_PREPARING] is export.ExportProgress.PREPARING
+
+
+def test_s3_settings_builders():
+    s = (
+        export.ExportToS3Settings()
+        .with_bucket("bucket")
+        .with_endpoint("endpoint")
+        .with_access_key("ak")
+        .with_secret_key("sk")
+        .with_scheme(1)
+        .with_uid("uid-1")
+        .with_number_of_retries(5)
+        .with_storage_class(2)
+        .with_export_compression("zstd")
+        .with_source_and_destination("/table", "prefix")
+        .with_items(("/t2", "p2"))
+        .with_item(("/t3", "p3"))
+    )
+    assert s.bucket == "bucket"
+    assert s.endpoint == "endpoint"
+    assert s.access_key == "ak"
+    assert s.secret_key == "sk"
+    assert s.scheme == 1
+    assert s.uid == "uid-1"
+    assert s.number_of_retries == 5
+    assert s.storage_class == 2
+    assert s.export_compression == "zstd"
+    assert s.items == [("/table", "prefix"), ("/t2", "p2"), ("/t3", "p3")]
+
+
+def test_s3_request_factory_defaults():
+    request = export._export_to_s3_request_factory(export.ExportToS3Settings())
+    assert request.settings.scheme == 2
+    assert len(request.settings.items) == 0
+    assert "uid" not in request.operation_params.labels
+    assert request.settings.number_of_retries == 0
+
+
+def test_s3_request_factory_full():
+    settings = (
+        export.ExportToS3Settings()
+        .with_bucket("b")
+        .with_endpoint("e")
+        .with_access_key("ak")
+        .with_secret_key("sk")
+        .with_uid("uid-9")
+        .with_number_of_retries(3)
+        .with_export_compression("gzip")
+        .with_source_and_destination("/table", "prefix")
+    )
+    request = export._export_to_s3_request_factory(settings)
+    assert request.settings.bucket == "b"
+    assert request.settings.endpoint == "e"
+    assert request.operation_params.labels["uid"] == "uid-9"
+    assert request.settings.number_of_retries == 3
+    assert request.settings.compression == "gzip"
+    assert request.settings.items[0].source_path == "/table"
+    assert request.settings.items[0].destination_prefix == "prefix"
+
+
+def test_yt_settings_builders():
+    s = (
+        export.ExportToYTSettings()
+        .with_host("host")
+        .with_port(8080)
+        .with_uid("uid")
+        .with_token("token")
+        .with_number_of_retries(2)
+        .with_use_type_v3(True)
+        .with_source_and_destination("/a", "/b")
+        .with_items(("/c", "/d"))
+    )
+    assert s.host == "host"
+    assert s.port == 8080
+    assert s.uid == "uid"
+    assert s.token == "token"
+    assert s.number_of_retries == 2
+    assert s.use_type_v3 is True
+    assert s.items == [("/a", "/b"), ("/c", "/d")]
+
+
+def test_yt_request_factory_defaults():
+    request = export._export_to_yt_request_factory(export.ExportToYTSettings())
+    assert request.settings.number_of_retries == 0
+    assert request.settings.port == 0
+    assert request.settings.use_type_v3 is False
+    assert len(request.settings.items) == 0
+
+
+def test_yt_request_factory_full():
+    settings = (
+        export.ExportToYTSettings()
+        .with_host("host")
+        .with_token("token")
+        .with_port(9090)
+        .with_number_of_retries(4)
+        .with_use_type_v3(True)
+        .with_source_and_destination("/a", "/b")
+    )
+    request = export._export_to_yt_request_factory(settings)
+    assert request.settings.host == "host"
+    assert request.settings.token == "token"
+    assert request.settings.port == 9090
+    assert request.settings.number_of_retries == 4
+    assert request.settings.use_type_v3 is True
+    assert request.settings.items[0].source_path == "/a"
+    assert request.settings.items[0].destination_path == "/b"
+
+
+def test_export_operation_wrappers():
+    s3 = export.ExportToS3Operation(None, _s3_response(), driver=None)
+    assert s3.id == "op-1"
+    assert s3.progress is export.ExportProgress.DONE
+    assert "ExportToS3Operation" in str(s3)
+    assert repr(s3) == str(s3)
+
+    yt = export.ExportToYTOperation(None, _yt_response(), driver=None)
+    assert yt.id == "op-2"
+    assert yt.progress is export.ExportProgress.TRANSFER_DATA
+    assert "ExportToYTOperation" in str(yt)
+    assert repr(yt) == str(yt)
+
+
+def test_export_client_dispatch():
+    driver = _FakeDriver()
+    client = export.ExportClient(driver)
+
+    client.export_to_s3(export.ExportToS3Settings())
+    request, stub, method, wrapper, settings, wrap_args = driver.calls[-1]
+    assert stub is ydb_export_v1_pb2_grpc.ExportServiceStub
+    assert method == export._ExportToS3
+    assert wrapper is export.ExportToS3Operation
+    assert wrap_args == (driver,)
+
+    client.export_to_yt(export.ExportToYTSettings())
+    _, _, method, wrapper, _, _ = driver.calls[-1]
+    assert method == export._ExportToYt
+    assert wrapper is export.ExportToYTOperation
+
+    client.async_export_to_yt(export.ExportToYTSettings())
+    _, _, method, wrapper, _, _ = driver.calls[-1]
+    assert method == export._ExportToYt
+
+    client.get_export_to_s3_operation("op-x")
+    request, _, method, wrapper, _, _ = driver.calls[-1]
+    assert request.id == "op-x"
+    assert method == _apis.OperationService.GetOperation
+    assert wrapper is export.ExportToS3Operation
+
+
+def test_export_operation_error_status_raises():
+    op = _apis.ydb_operation.Operation(status=issues.StatusCode.BAD_REQUEST)
+    op.metadata.Pack(ydb_export_pb2.ExportToS3Metadata())
+    response = _apis.ydb_operation.GetOperationResponse(operation=op)
+    with pytest.raises(issues.BadRequest):
+        export.ExportToS3Operation(None, response, driver=None)