def test_create_dask_cloud_provider_environment_aws_creds_provided(): environment = DaskCloudProviderEnvironment( provider_class=FargateCluster, labels=["foo"], aws_access_key_id="id", aws_secret_access_key="secret", aws_session_token="session", region_name="region", ) assert environment assert environment.labels == set(["foo"]) assert environment._provider_kwargs["aws_access_key_id"] == "id" assert environment._provider_kwargs["aws_secret_access_key"] == "secret" assert environment._provider_kwargs["aws_session_token"] == "session" assert environment._provider_kwargs["region_name"] == "region"
def test_create_dask_cloud_provider_environment_callbacks(): def f(): pass environment = DaskCloudProviderEnvironment( provider_class=FargateCluster, labels=["foo"], on_execute=f, on_start=f, on_exit=f, ) assert environment assert environment.labels == set(["foo"]) assert environment._on_execute is f assert environment.on_start is f assert environment.on_exit is f
def test_roundtrip_cloudpickle(): with tempfile.TemporaryDirectory() as directory: with open(os.path.join(directory, "job.yaml"), "w+") as file: file.write("job") environment = DaskCloudProviderEnvironment( provider_class=FargateCluster, cluster="test", skip_cleanup=False, ) assert environment._provider_kwargs == { "cluster": "test", "skip_cleanup": False, } new = cloudpickle.loads(cloudpickle.dumps(environment)) assert isinstance(new, DaskCloudProviderEnvironment) assert new._provider_kwargs == {"cluster": "test", "skip_cleanup": False}
def test_dask_cloud_provider_environment_dependencies(): environment = DaskCloudProviderEnvironment(provider_class=FargateCluster) assert environment.dependencies == ["dask_cloudprovider"]
def test_create_dask_cloud_provider_environment_labels(): environment = DaskCloudProviderEnvironment(provider_class=FargateCluster, labels=["foo"]) assert environment assert environment.labels == set(["foo"])
def test_create_dask_cloud_provider_environment_with_executor_kwargs(): environment = DaskCloudProviderEnvironment( provider_class=FargateCluster, executor_kwargs={"test": "here"}) assert environment assert environment.executor_kwargs == {"address": "", "test": "here"}
def test_create_environment(): environment = DaskCloudProviderEnvironment(Cluster) assert environment