download_config_dict init fix

Couchbase-Ecosystem · Dec 10, 2024 · 055a814 · 055a814
1 parent e3c4f1c
commit 055a814
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git 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
@@ -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)