[GitHub] [beam] robertwb commented on a change in pull request #11305: [BEAM-9577] Update container boot code to stage from dependencies, if present.

[GitBox <gi...@apache.org>]

robertwb commented on a change in pull request #11305: [BEAM-9577] Update container boot code to stage from dependencies, if present.
URL: https://github.com/apache/beam/pull/11305#discussion_r403281683
 
 

 ##########
 File path: sdks/go/pkg/beam/artifact/materialize_test.go
 ##########
 @@ -148,6 +153,156 @@ func stage(ctx context.Context, scl pb.LegacyArtifactStagingServiceClient, t *te
 	return md
 }
 
+// Test for new artifact retrieval.
+
+func TestNewRetrieveWithManyFiles(t *testing.T) {
+	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}
+
+	client := fakeRetrievalService{
+		artifacts: expected,
+	}
+
+	dest := makeTempDir(t)
+	defer os.RemoveAll(dest)
+	ctx := grpcx.WriteWorkerID(context.Background(), "worker")
+
+	mds, err := newMaterializeWithClient(ctx, client, client.resolvedArtifacts(), dest)
+	if err != nil {
+		t.Fatalf("materialize failed: %v", err)
+	}
+
+	checkStagedFiles(mds, dest, expected, t)
+}
+
+func TestNewRetrieveWithResolution(t *testing.T) {
+	expected := map[string]string{"a.txt": "a", "b.txt": "bbb", "c.txt": "cccccccc"}
+
+	client := fakeRetrievalService{
+		artifacts: expected,
+	}
+
+	dest := makeTempDir(t)
+	defer os.RemoveAll(dest)
+	ctx := grpcx.WriteWorkerID(context.Background(), "worker")
+
+	mds, err := newMaterializeWithClient(ctx, client, client.unresolvedArtifacts(), dest)
+	if err != nil {
+		t.Fatalf("materialize failed: %v", err)
+	}
+
+	checkStagedFiles(mds, dest, expected, t)
+}
+
+func checkStagedFiles(mds []*pb.ArtifactMetadata, dest string, expected map[string]string, t *testing.T) {
+	if len(mds) != len(expected) {
+		t.Errorf("wrong number of artifacts staged %v vs %v", len(mds), len(expected))
+	}
+	for _, md := range mds {
+		filename := filepath.Join(dest, filepath.FromSlash(md.Name))
+		fd, err := os.Open(filename)
+		if err != nil {
+			t.Errorf("error opening file %v", err)
+		}
+		defer fd.Close()
+
+		data := make([]byte, 1<<20)
+		n, err := fd.Read(data)
+		if err != nil {
+			t.Errorf("error reading file %v", err)
+		}
+
+		if string(data[:n]) != expected[md.Name] {
+			t.Errorf("missmatched contents for %v: '%s' vs '%s'", md.Name, string(data[:n]), expected[md.Name])
+		}
+	}
+}
+
+type fakeRetrievalService struct {
+	artifacts map[string]string // name -> content
+}
+
+func (fake fakeRetrievalService) resolvedArtifacts() []*pipeline_v1.ArtifactInformation {
+	var artifacts []*pipeline_v1.ArtifactInformation
+	for name, contents := range fake.artifacts {
+		payload, _ := proto.Marshal(&pipeline_v1.ArtifactStagingToRolePayload{
+			StagedName: name})
+		artifacts = append(artifacts, &pipeline_v1.ArtifactInformation{
+			TypeUrn:     "resolved",
+			TypePayload: []byte(contents),
+			RoleUrn:     URNStagingTo,
+			RolePayload: payload,
+		})
+	}
+	return artifacts
+}
+
+func (fake fakeRetrievalService) unresolvedArtifacts() []*pipeline_v1.ArtifactInformation {
+	return []*pipeline_v1.ArtifactInformation{
+		&pipeline_v1.ArtifactInformation{
+			TypeUrn: "unresolved",
+		},
+	}
+}
+
+func (fake fakeRetrievalService) ResolveArtifact(ctx context.Context, request *pb.ResolveArtifactRequest, opts ...grpc.CallOption) (*pb.ResolveArtifactResponse, error) {
+	response := pb.ResolveArtifactResponse{}
+	for _, dep := range request.Artifacts {
+		if dep.TypeUrn == "unresolved" {
+			response.Replacements = append(response.Replacements, fake.resolvedArtifacts()...)
+		} else {
+			response.Replacements = append(response.Replacements, dep)
+		}
+	}
+	return &response, nil
+}
+
+func (fake fakeRetrievalService) GetArtifact(ctx context.Context, request *pb.GetArtifactRequest, opts ...grpc.CallOption) (pb.ArtifactRetrievalService_GetArtifactClient, error) {
+	var index int
+	if request.Artifact.TypeUrn == "resolved" {
+		return fakeGetArtifactResponse{data: request.Artifact.TypePayload, index: &index}, nil
+	} else {
+		return nil, errors.Errorf("Unsupported artifact %v", request.Artifact)
+	}
+}
+
+func (fake fakeGetArtifactResponse) Recv() (*pb.GetArtifactResponse, error) {
+	if *fake.index < len(fake.data) {
+		*fake.index += 1
+		return &pb.GetArtifactResponse{Data: fake.data[*fake.index-1 : *fake.index]}, nil
+	} else {
+		return nil, io.EOF
+	}
+}
+
+type fakeGetArtifactResponse struct {
+	data  []byte
+	index *int
 
 Review comment:
   That is what I tried at first, but it complained that it didn't satisfy the interface. 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services