github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/portability/artifact_service_test.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Test cases for :module:`artifact_service_client`."""

# pytype: skip-file

import contextlib
import io
import threading
import unittest
from urllib.parse import quote

from apache_beam.portability import common_urns
from apache_beam.portability.api import beam_artifact_api_pb2
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.runners.portability import artifact_service
from apache_beam.utils import proto_utils


class InMemoryFileManager(object):
  def __init__(self, contents=()):
    self._contents = dict(contents)

  def get(self, path):
    return self._contents[path]

  def file_reader(self, path):
    return io.BytesIO(self._contents[path])

  def file_writer(self, name):
    path = 'prefix:' + name

    @contextlib.contextmanager
    def writable():
      buffer = io.BytesIO()
      yield buffer
      buffer.seek(0)
      self._contents[path] = buffer.read()

    return writable(), path


class ArtifactServiceTest(unittest.TestCase):
  def file_artifact(self, path):
    return beam_runner_api_pb2.ArtifactInformation(
        type_urn=common_urns.artifact_types.FILE.urn,
        type_payload=beam_runner_api_pb2.ArtifactFilePayload(
            path=path).SerializeToString())

  def embedded_artifact(self, data, name=None):
    return beam_runner_api_pb2.ArtifactInformation(
        type_urn=common_urns.artifact_types.EMBEDDED.urn,
        type_payload=beam_runner_api_pb2.EmbeddedFilePayload(
            data=data).SerializeToString(),
        role_urn=common_urns.artifact_roles.STAGING_TO.urn if name else None,
        role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload(
            staged_name=name).SerializeToString() if name else None)

  def test_file_retrieval(self):
    file_manager = InMemoryFileManager({
        'path/to/a': b'a', 'path/to/b': b'b' * 37
    })
    retrieval_service = artifact_service.ArtifactRetrievalService(
        file_manager.file_reader, chunk_size=10)
    dep_a = self.file_artifact('path/to/a')
    self.assertEqual(
        retrieval_service.ResolveArtifacts(
            beam_artifact_api_pb2.ResolveArtifactsRequest(artifacts=[dep_a])),
        beam_artifact_api_pb2.ResolveArtifactsResponse(replacements=[dep_a]))

    self.assertEqual(
        list(
            retrieval_service.GetArtifact(
                beam_artifact_api_pb2.GetArtifactRequest(artifact=dep_a))),
        [beam_artifact_api_pb2.GetArtifactResponse(data=b'a')])
    self.assertEqual(
        list(
            retrieval_service.GetArtifact(
                beam_artifact_api_pb2.GetArtifactRequest(
                    artifact=self.file_artifact('path/to/b')))),
        [
            beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 10),
            beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 10),
            beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 10),
            beam_artifact_api_pb2.GetArtifactResponse(data=b'b' * 7)
        ])

  def test_embedded_retrieval(self):
    retrieval_service = artifact_service.ArtifactRetrievalService(None)
    embedded_dep = self.embedded_artifact(b'some_data')
    self.assertEqual(
        list(
            retrieval_service.GetArtifact(
                beam_artifact_api_pb2.GetArtifactRequest(
                    artifact=embedded_dep))),
        [beam_artifact_api_pb2.GetArtifactResponse(data=b'some_data')])

  def test_url_retrieval(self):
    retrieval_service = artifact_service.ArtifactRetrievalService(None)
    url_dep = beam_runner_api_pb2.ArtifactInformation(
        type_urn=common_urns.artifact_types.URL.urn,
        type_payload=beam_runner_api_pb2.ArtifactUrlPayload(
            url='file:' + quote(__file__)).SerializeToString())
    content = b''.join([
        r.data for r in retrieval_service.GetArtifact(
            beam_artifact_api_pb2.GetArtifactRequest(artifact=url_dep))
    ])
    with open(__file__, 'rb') as fin:
      self.assertEqual(content, fin.read())

  def test_push_artifacts(self):
    unresolved = beam_runner_api_pb2.ArtifactInformation(type_urn='unresolved')
    resolved_a = self.embedded_artifact(data=b'a', name='a.txt')
    resolved_b = self.embedded_artifact(data=b'bb', name='b.txt')
    dep_big = self.embedded_artifact(data=b'big ' * 100, name='big.txt')

    class TestArtifacts(object):
      def ResolveArtifacts(self, request):
        replacements = []
        for artifact in request.artifacts:
          if artifact.type_urn == 'unresolved':
            replacements += [resolved_a, resolved_b]
          else:
            replacements.append(artifact)
        return beam_artifact_api_pb2.ResolveArtifactsResponse(
            replacements=replacements)

      def GetArtifact(self, request):
        if request.artifact.type_urn == common_urns.artifact_types.EMBEDDED.urn:
          content = proto_utils.parse_Bytes(
              request.artifact.type_payload,
              beam_runner_api_pb2.EmbeddedFilePayload).data
          for k in range(0, len(content), 13):
            yield beam_artifact_api_pb2.GetArtifactResponse(
                data=content[k:k + 13])
        else:
          raise NotImplementedError

    file_manager = InMemoryFileManager()
    server = artifact_service.ArtifactStagingService(file_manager.file_writer)

    server.register_job('staging_token', {'env': [unresolved, dep_big]})

    # "Push" artifacts as if from a client.
    t = threading.Thread(
        target=lambda: artifact_service.offer_artifacts(
            server, TestArtifacts(), 'staging_token'))
    t.daemon = True
    t.start()

    resolved_deps = server.resolved_deps('staging_token', timeout=5)['env']
    expected = {
        'a.txt': b'a',
        'b.txt': b'bb',
        'big.txt': b'big ' * 100,
    }
    for dep in resolved_deps:
      self.assertEqual(dep.type_urn, common_urns.artifact_types.FILE.urn)
      self.assertEqual(dep.role_urn, common_urns.artifact_roles.STAGING_TO.urn)
      type_payload = proto_utils.parse_Bytes(
          dep.type_payload, beam_runner_api_pb2.ArtifactFilePayload)
      role_payload = proto_utils.parse_Bytes(
          dep.role_payload, beam_runner_api_pb2.ArtifactStagingToRolePayload)
      self.assertTrue(
          type_payload.path.endswith(role_payload.staged_name),
          type_payload.path)
      self.assertEqual(
          file_manager.get(type_payload.path),
          expected.pop(role_payload.staged_name))
    self.assertEqual(expected, {})


if __name__ == '__main__':
  unittest.main()