From 055a814c1ccab80b4126a3dbd77cab1940658bc6 Mon Sep 17 00:00:00 2001 From: Shyam Venkat Date: Tue, 10 Dec 2024 17:07:48 +0530 Subject: [PATCH] download_config_dict init fix --- hf_to_cb_dataset_migrator/cli.py | 8 ++++++-- tests/integration/test_integration.py | 10 +++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/hf_to_cb_dataset_migrator/cli.py b/hf_to_cb_dataset_migrator/cli.py index 793f39b..40b1d8e 100644 --- a/hf_to_cb_dataset_migrator/cli.py +++ b/hf_to_cb_dataset_migrator/cli.py @@ -64,14 +64,15 @@ def list_configs_cmd(ctx, path, revision, download_config, download_mode, dynami logging.basicConfig(level=logging.DEBUG) migrator = DatasetMigrator(token=token) + + download_config_dict = None if download_config: try: download_config_dict = json.loads(download_config) except json.JSONDecodeError as e: click.echo(f"Error parsing download_config JSON: {e}", err=True) sys.exit(1) - else: - download_config_dict = None + download_kwargs = { 'revision': revision, @@ -137,6 +138,7 @@ def list_splits_cmd(ctx, path, config_name, data_files, download_config, downloa elif debug: logging.basicConfig(level=logging.DEBUG) + download_config_dict = None if download_config: try: download_config_dict = json.loads(download_config) @@ -207,6 +209,7 @@ def list_fields(ctx, path, name, data_files, download_config, revision, token, s elif debug: logging.basicConfig(level=logging.DEBUG) + download_config_dict = None if download_config: try: download_config_dict = json.loads(download_config) @@ -333,6 +336,7 @@ def migrate( if debug: logging.basicConfig(level=logging.DEBUG) + download_config_dict = None if download_config: try: download_config_dict = json.loads(download_config) diff --git a/tests/integration/test_integration.py b/tests/integration/test_integration.py index bd02445..a9bc5d6 100644 --- a/tests/integration/test_integration.py +++ b/tests/integration/test_integration.py @@ -10,7 +10,7 @@ import logging from couchbase.cluster import Cluster from couchbase.auth import PasswordAuthenticator -from couchbase.options import ClusterOptions +from couchbase.options import ClusterOptions, KnownConfigProfiles from datasets import Dataset, DatasetDict from tempfile import TemporaryDirectory @@ -67,7 +67,9 @@ def cleanup_collection(): """Fixture to clean up test collection before and after each test""" def _cleanup_collection(scope_name: str, collection_name: str): auth = PasswordAuthenticator(COUCHBASE_USERNAME, COUCHBASE_PASSWORD) - cluster = Cluster(COUCHBASE_URL, ClusterOptions(auth)) + cluster_opts = ClusterOptions(auth) + cluster_opts.apply_profile(KnownConfigProfiles.WanDevelopment) + cluster = Cluster(COUCHBASE_URL, cluster_opts) bucket = cluster.bucket(COUCHBASE_BUCKET) try: @@ -416,7 +418,9 @@ def validate_migrated_data( try: # Connect to Couchbase auth = PasswordAuthenticator(cb_username, cb_password) - cluster = Cluster(cb_url, ClusterOptions(auth)) + cluster_opts = ClusterOptions(auth) + cluster_opts.apply_profile(KnownConfigProfiles.WanDevelopment) + cluster = Cluster(cb_url, cluster_opts) bucket = cluster.bucket(cb_bucket) scope = bucket.scope(cb_scope) collection = scope.collection(cb_collection)