|
2 | 2 | from __future__ import annotations
|
3 | 3 |
|
4 | 4 | from pathlib import Path
|
5 |
| -from typing import Optional |
| 5 | +from typing import Optional, Union |
6 | 6 | from unittest import mock
|
7 | 7 | from unittest.mock import MagicMock
|
8 | 8 |
|
@@ -416,13 +416,13 @@ def test_delete_declarative_data_sources(test_config):
|
416 | 416 | credentials_path = _current_dir / "load" / "data_source_credentials" / "data_sources_credentials.yaml"
|
417 | 417 | expected_json_path = _current_dir / "expected" / "declarative_data_sources.json"
|
418 | 418 |
|
419 |
| - def token_from_file_side_effect(arg): |
420 |
| - if arg == "~/home/secrets.json": |
| 419 | + def token_from_file_side_effect(file_path: Union[str, Path], base64_encode: bool): |
| 420 | + if file_path == "~/home/secrets.json": |
421 | 421 | return test_config["bigquery_token"]
|
422 |
| - elif arg == "databricks-token": |
| 422 | + elif file_path == "databricks-token": |
423 | 423 | return test_config["databricks_token"]
|
424 | 424 | else:
|
425 |
| - raise ValueError(f"Unexpected argument: {arg}") |
| 425 | + raise ValueError(f"Unexpected argument: {file_path}") |
426 | 426 |
|
427 | 427 | TokenCredentialsFromFile.token_from_file = MagicMock(side_effect=token_from_file_side_effect)
|
428 | 428 | ClientSecretCredentialsFromFile.client_secret_from_file = MagicMock(
|
@@ -460,13 +460,13 @@ def test_load_and_put_declarative_data_sources(test_config):
|
460 | 460 | try:
|
461 | 461 | sdk.catalog_data_source.put_declarative_data_sources(CatalogDeclarativeDataSources(data_sources=[]))
|
462 | 462 |
|
463 |
| - def token_from_file_side_effect(arg): |
464 |
| - if arg == "~/home/secrets.json": |
| 463 | + def token_from_file_side_effect(file_path: Union[str, Path], base64_encode: bool = True): |
| 464 | + if file_path == "~/home/secrets.json": |
465 | 465 | return test_config["bigquery_token"]
|
466 |
| - elif arg == "databricks-token": |
| 466 | + elif file_path == "databricks-token": |
467 | 467 | return test_config["databricks_token"]
|
468 | 468 | else:
|
469 |
| - raise ValueError(f"Unexpected argument: {arg}") |
| 469 | + raise ValueError(f"Unexpected argument: {file_path}") |
470 | 470 |
|
471 | 471 | TokenCredentialsFromFile.token_from_file = MagicMock(side_effect=token_from_file_side_effect)
|
472 | 472 | ClientSecretCredentialsFromFile.client_secret_from_file = MagicMock(
|
@@ -506,13 +506,13 @@ def test_put_declarative_data_sources_connection(test_config):
|
506 | 506 | # Must filter out databricks data sources for this test as they do not have valid URLs
|
507 | 507 | data_sources_e.data_sources = [item for item in data_sources_e.data_sources if "databricks" not in item.id]
|
508 | 508 |
|
509 |
| - def token_from_file_side_effect(arg): |
510 |
| - if arg == "~/home/secrets.json": |
| 509 | + def token_from_file_side_effect(file_path: Union[str, Path], base64_encode: bool): |
| 510 | + if file_path == "~/home/secrets.json": |
511 | 511 | return test_config["bigquery_token"]
|
512 |
| - elif arg == "databricks-token": |
| 512 | + elif file_path == "databricks-token": |
513 | 513 | return test_config["databricks_token"]
|
514 | 514 | else:
|
515 |
| - raise ValueError(f"Unexpected argument: {arg}") |
| 515 | + raise ValueError(f"Unexpected argument: {file_path}") |
516 | 516 |
|
517 | 517 | TokenCredentialsFromFile.token_from_file = MagicMock(side_effect=token_from_file_side_effect)
|
518 | 518 | ClientSecretCredentialsFromFile.client_secret_from_file = MagicMock(
|
@@ -605,13 +605,13 @@ def test_cache_strategy(test_config: dict):
|
605 | 605 | path = _current_dir / "expected" / "declarative_data_sources.json"
|
606 | 606 | credentials_path = _current_dir / "load" / "data_source_credentials" / "data_sources_credentials.yaml"
|
607 | 607 |
|
608 |
| - def token_from_file_side_effect(arg): |
609 |
| - if arg == "~/home/secrets.json": |
| 608 | + def token_from_file_side_effect(file_path: Union[str, Path], base64_encode: bool): |
| 609 | + if file_path == "~/home/secrets.json": |
610 | 610 | return test_config["bigquery_token"]
|
611 |
| - elif arg == "databricks-token": |
| 611 | + elif file_path == "databricks-token": |
612 | 612 | return test_config["databricks_token"]
|
613 | 613 | else:
|
614 |
| - raise ValueError(f"Unexpected argument: {arg}") |
| 614 | + raise ValueError(f"Unexpected argument: {file_path}") |
615 | 615 |
|
616 | 616 | TokenCredentialsFromFile.token_from_file = MagicMock(side_effect=token_from_file_side_effect)
|
617 | 617 | ClientSecretCredentialsFromFile.client_secret_from_file = MagicMock(
|
|
0 commit comments