@pytest.mark.parametrize('labels', [[], ['b', 'c']])
@pytest.mark.parametrize('kind', ['run_config', 'environment'])
def test_register_flow_call(monkeypatch, tmpdir, kind, labels):
    client = MagicMock()
    monkeypatch.setattr('prefect.Client', MagicMock(return_value=client))
    if (kind == 'environment'):
        contents = "from prefect import Flow\nfrom prefect.environments.execution import LocalEnvironment\nfrom prefect.environments.storage import Local\nf = Flow('test-flow', environment=LocalEnvironment(labels=['a']),\n   storage=Local(add_default_labels=False))"
    else:
        contents = "from prefect import Flow\nfrom prefect.run_configs import KubernetesRun\nfrom prefect.environments.storage import Local\nf = Flow('test-flow', run_config=KubernetesRun(labels=['a']),\n   storage=Local(add_default_labels=False))"
    full_path = str(tmpdir.join('flow.py'))
    with open(full_path, 'w') as f:
        f.write(contents)
    args = ['flow', '--file', full_path, '--name', 'test-flow', '--project', 'project']
    for l in labels:
        args.extend(['-l', l])
    runner = CliRunner()
    result = runner.invoke(register, args)
    assert client.register.called
    assert (client.register.call_args[1]['project_name'] == 'project')
    flow = client.register.call_args[1]['flow']
    if (kind == 'run_config'):
        assert (flow.run_config.labels == {'a', *labels})
    else:
        assert (flow.environment.labels == {'a', *labels})
    assert (result.exit_code == 0)
@pytest.mark.parametrize('labels', [[], ['b', 'c']])
@pytest.mark.parametrize('kind', ['run_config', 'environment'])
def test_register_flow_call(monkeypatch, tmpdir, kind, labels):
    client = MagicMock()
    monkeypatch.setattr('prefect.Client', MagicMock(return_value=client))
    if (kind == 'environment'):
        contents = "from prefect import Flow\nfrom prefect.environments.execution import LocalEnvironment\nfrom prefect.environments.storage import Local\nf = Flow('test-flow', environment=LocalEnvironment(labels=['a']),\n   storage=Local(add_default_labels=False))"
    else:
        contents = "from prefect import Flow\nfrom prefect.run_configs import KubernetesRun\nfrom prefect.environments.storage import Local\nf = Flow('test-flow', run_config=KubernetesRun(labels=['a']),\n   storage=Local(add_default_labels=False))"
    full_path = str(tmpdir.join('flow.py'))
    with open(full_path, 'w') as f:
        f.write(contents)
    args = ['flow', '--file', full_path, '--name', 'test-flow', '--project', 'project']
    for l in labels:
        args.extend(['-l', l])
    runner = CliRunner()
    result = runner.invoke(register, args)
    assert client.register.called
    assert (client.register.call_args[1]['project_name'] == 'project')
    flow = client.register.call_args[1]['flow']
    if (kind == 'run_config'):
        assert (flow.run_config.labels == {'a', *labels})
    else:
        assert (flow.environment.labels == {'a', *labels})
    assert (result.exit_code == 0)