def test__create_default_environment(self):
docker_image = PortableRunner.default_docker_image()
self.assertEqual(
PortableRunner._create_environment(PipelineOptions.from_dictionary({})),
beam_runner_api_pb2.Environment(
urn=common_urns.environments.DOCKER.urn,
payload=beam_runner_api_pb2.DockerPayload(
container_image=docker_image
).SerializeToString()))
def test__create_default_environment(self):
docker_image = environments.DockerEnvironment.default_docker_image()
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({'sdk_location':
'container'})),
environments.DockerEnvironment(container_image=docker_image))
def __init__(self, container_image=None):
from apache_beam.runners.portability.portable_runner import PortableRunner
if container_image:
self.container_image = container_image
else:
self.container_image = PortableRunner.default_docker_image()
def hasDockerImage():
image = PortableRunner.default_docker_image()
try:
check_image = subprocess.check_output("docker images -q %s" % image,
shell=True)
return check_image != ''
except Exception:
return False
def test__create_process_environment(self):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': "PROCESS",
'environment_config': '{"os": "linux", "arch": "amd64", '
'"command": "run.sh", '
'"env":{"k1": "v1"} }',
})),
environments.ProcessEnvironment(
'run.sh', os='linux', arch='amd64', env={'k1': 'v1'}))
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': 'PROCESS',
'environment_config': '{"command": "run.sh"}',
})),
environments.ProcessEnvironment('run.sh'))
def test__create_docker_environment(self):
docker_image = 'py-docker'
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': 'DOCKER',
'environment_config': docker_image,
})),
environments.DockerEnvironment(container_image=docker_image))
def test__create_docker_environment(self):
docker_image = 'py-docker'
self.assertEqual(
PortableRunner._create_environment(PipelineOptions.from_dictionary({
'environment_type': 'DOCKER',
'environment_config': docker_image,
})), beam_runner_api_pb2.Environment(
urn=common_urns.environments.DOCKER.urn,
payload=beam_runner_api_pb2.DockerPayload(
container_image=docker_image
).SerializeToString()))
def test__create_external_environment(self):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': "EXTERNAL",
'environment_config': 'localhost:50000',
})),
environments.ExternalEnvironment('localhost:50000'))
raw_config = ' {"url":"localhost:50000", "params":{"k1":"v1"}} '
for env_config in (raw_config, raw_config.lstrip(), raw_config.strip()):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': "EXTERNAL",
'environment_config': env_config,
})),
environments.ExternalEnvironment(
'localhost:50000', params={"k1": "v1"}))
with self.assertRaises(ValueError):
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': "EXTERNAL",
'environment_config': '{invalid}',
}))
with self.assertRaises(ValueError) as ctx:
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type': "EXTERNAL",
'environment_config': '{"params":{"k1":"v1"}}',
}))
self.assertIn(
'External environment endpoint must be set.', ctx.exception.args)
def test__create_process_environment(self):
self.assertEqual(
PortableRunner._create_environment(PipelineOptions.from_dictionary({
'environment_type': "PROCESS",
'environment_config': '{"os": "linux", "arch": "amd64", '
'"command": "run.sh", '
'"env":{"k1": "v1"} }',
})), beam_runner_api_pb2.Environment(
urn=common_urns.environments.PROCESS.urn,
payload=beam_runner_api_pb2.ProcessPayload(
os='linux',
arch='amd64',
command='run.sh',
env={'k1': 'v1'},
).SerializeToString()))
self.assertEqual(
PortableRunner._create_environment(PipelineOptions.from_dictionary({
'environment_type': 'PROCESS',
'environment_config': '{"command": "run.sh"}',
})), beam_runner_api_pb2.Environment(
urn=common_urns.environments.PROCESS.urn,
payload=beam_runner_api_pb2.ProcessPayload(
command='run.sh',
).SerializeToString()))
def test__create_process_environment(self):
self.assertEqual(
PortableRunner._create_environment(PipelineOptions.from_dictionary({
'environment_type': "PROCESS",
'environment_config': '{"os": "linux", "arch": "amd64", '
'"command": "run.sh", '
'"env":{"k1": "v1"} }',
})), beam_runner_api_pb2.Environment(
urn=common_urns.environments.PROCESS.urn,
payload=beam_runner_api_pb2.ProcessPayload(
os='linux',
arch='amd64',
command='run.sh',
env={'k1': 'v1'},
).SerializeToString()))
self.assertEqual(
PortableRunner._create_environment(PipelineOptions.from_dictionary({
'environment_type': 'PROCESS',
'environment_config': '{"command": "run.sh"}',
})), beam_runner_api_pb2.Environment(
urn=common_urns.environments.PROCESS.urn,
payload=beam_runner_api_pb2.ProcessPayload(
command='run.sh',
).SerializeToString()))
def test__create_external_environment(self):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
'localhost:50000',
})),
beam_runner_api_pb2.Environment(
urn=common_urns.environments.EXTERNAL.urn,
payload=beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(
url='localhost:50000')).SerializeToString()))
raw_config = ' {"url":"localhost:50000", "params":{"test":"test"}} '
for env_config in (raw_config, raw_config.lstrip(),
raw_config.strip()):
self.assertEqual(
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
env_config,
})),
beam_runner_api_pb2.Environment(
urn=common_urns.environments.EXTERNAL.urn,
payload=beam_runner_api_pb2.ExternalPayload(
endpoint=endpoints_pb2.ApiServiceDescriptor(
url='localhost:50000'),
params={
"test": "test"
}).SerializeToString()))
with self.assertRaises(ValueError):
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
'{invalid}',
}))
with self.assertRaises(ValueError) as ctx:
PortableRunner._create_environment(
PipelineOptions.from_dictionary({
'environment_type':
"EXTERNAL",
'environment_config':
'{"params":{"test":"test"}}',
}))
self.assertIn('External environment endpoint must be set.',
ctx.exception.args)
def hasDockerImage():
image = PortableRunner.default_docker_image()
check_image = subprocess.check_output(["docker", "images", "-q", image])
return check_image != ''
