from dagster_databricks import DatabricksRunJobSolidDefinition

from dagster import execute_pipeline, pipeline
from dagster.seven import mock

HOST = 'https://uksouth.azuredatabricks.net'
TOKEN = 'super-secret-token'


@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.submit_run')
@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.wait_for_run_to_complete')
def test_run_databricks_job(mock_wait_for_run, mock_submit_run, databricks_run_config):
    @pipeline
    def test_pipe():
        DatabricksRunJobSolidDefinition('test', poll_interval_sec=1)()

    RUN_ID = 1
    mock_submit_run.return_value = RUN_ID

    config = {
        'solids': {
            'test': {
                'config': {
                    'run_config': databricks_run_config,
                    'databricks_host': HOST,
                    'databricks_token': TOKEN,
                }
            }
        }
    }
    result = execute_pipeline(test_pipe, config)
    assert result.success

    assert mock_submit_run.call_count == 1
    task = databricks_run_config.pop('task')
    assert mock_submit_run.call_args[0][0] == databricks_run_config
    assert mock_submit_run.call_args[0][1] == task
    assert mock_wait_for_run.call_count == 1
    assert mock_wait_for_run.call_args[0][1] == RUN_ID

af14a4ea9f92dd93b6f9c026cdcde933c2eb670c
<mask token>

<mask token>


@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.submit_run')
@mock.patch(
    'dagster_databricks.databricks.DatabricksJobRunner.wait_for_run_to_complete'
    )
def test_run_databricks_job(mock_wait_for_run, mock_submit_run,
    databricks_run_config):

    @pipeline
    def test_pipe():
        DatabricksRunJobSolidDefinition('test', poll_interval_sec=1)()
    RUN_ID = 1
    mock_submit_run.return_value = RUN_ID
    config = {'solids': {'test': {'config': {'run_config':
        databricks_run_config, 'databricks_host': HOST, 'databricks_token':
        TOKEN}}}}
    result = execute_pipeline(test_pipe, config)
    assert result.success
    assert mock_submit_run.call_count == 1
    task = databricks_run_config.pop('task')
    assert mock_submit_run.call_args[0][0] == databricks_run_config
    assert mock_submit_run.call_args[0][1] == task
    assert mock_wait_for_run.call_count == 1
    assert mock_wait_for_run.call_args[0][1] == RUN_ID

<mask token>
HOST = 'https://uksouth.azuredatabricks.net'
TOKEN = 'super-secret-token'


@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.submit_run')
@mock.patch(
    'dagster_databricks.databricks.DatabricksJobRunner.wait_for_run_to_complete'
    )
def test_run_databricks_job(mock_wait_for_run, mock_submit_run,
    databricks_run_config):

    @pipeline
    def test_pipe():
        DatabricksRunJobSolidDefinition('test', poll_interval_sec=1)()
    RUN_ID = 1
    mock_submit_run.return_value = RUN_ID
    config = {'solids': {'test': {'config': {'run_config':
        databricks_run_config, 'databricks_host': HOST, 'databricks_token':
        TOKEN}}}}
    result = execute_pipeline(test_pipe, config)
    assert result.success
    assert mock_submit_run.call_count == 1
    task = databricks_run_config.pop('task')
    assert mock_submit_run.call_args[0][0] == databricks_run_config
    assert mock_submit_run.call_args[0][1] == task
    assert mock_wait_for_run.call_count == 1
    assert mock_wait_for_run.call_args[0][1] == RUN_ID

from dagster_databricks import DatabricksRunJobSolidDefinition
from dagster import execute_pipeline, pipeline
from dagster.seven import mock
HOST = 'https://uksouth.azuredatabricks.net'
TOKEN = 'super-secret-token'


@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.submit_run')
@mock.patch(
    'dagster_databricks.databricks.DatabricksJobRunner.wait_for_run_to_complete'
    )
def test_run_databricks_job(mock_wait_for_run, mock_submit_run,
    databricks_run_config):

    @pipeline
    def test_pipe():
        DatabricksRunJobSolidDefinition('test', poll_interval_sec=1)()
    RUN_ID = 1
    mock_submit_run.return_value = RUN_ID
    config = {'solids': {'test': {'config': {'run_config':
        databricks_run_config, 'databricks_host': HOST, 'databricks_token':
        TOKEN}}}}
    result = execute_pipeline(test_pipe, config)
    assert result.success
    assert mock_submit_run.call_count == 1
    task = databricks_run_config.pop('task')
    assert mock_submit_run.call_args[0][0] == databricks_run_config
    assert mock_submit_run.call_args[0][1] == task
    assert mock_wait_for_run.call_count == 1
    assert mock_wait_for_run.call_args[0][1] == RUN_ID

from dagster_databricks import DatabricksRunJobSolidDefinition

from dagster import execute_pipeline, pipeline
from dagster.seven import mock

HOST = 'https://uksouth.azuredatabricks.net'
TOKEN = 'super-secret-token'


@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.submit_run')
@mock.patch('dagster_databricks.databricks.DatabricksJobRunner.wait_for_run_to_complete')
def test_run_databricks_job(mock_wait_for_run, mock_submit_run, databricks_run_config):
    @pipeline
    def test_pipe():
        DatabricksRunJobSolidDefinition('test', poll_interval_sec=1)()

    RUN_ID = 1
    mock_submit_run.return_value = RUN_ID

    config = {
        'solids': {
            'test': {
                'config': {
                    'run_config': databricks_run_config,
                    'databricks_host': HOST,
                    'databricks_token': TOKEN,
                }
            }
        }
    }
    result = execute_pipeline(test_pipe, config)
    assert result.success

    assert mock_submit_run.call_count == 1
    task = databricks_run_config.pop('task')
    assert mock_submit_run.call_args[0][0] == databricks_run_config
    assert mock_submit_run.call_args[0][1] == task
    assert mock_wait_for_run.call_count == 1
    assert mock_wait_for_run.call_args[0][1] == RUN_ID
