github.com/apache/beam/sdks/v2@v2.48.2/go/container/boot_test.go

// 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.

package main

import (
	"context"
	"os"
	"path/filepath"
	"testing"

	"github.com/apache/beam/sdks/v2/go/container/tools"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
	"github.com/golang/protobuf/proto"
)

func TestEnsureEndpointsSet_AllSet(t *testing.T) {
	provisionInfo := &fnpb.ProvisionInfo{
		LoggingEndpoint:  &pipepb.ApiServiceDescriptor{Url: "testLoggingEndpointUrl"},
		ArtifactEndpoint: &pipepb.ApiServiceDescriptor{Url: "testArtifactEndpointUrl"},
		ControlEndpoint:  &pipepb.ApiServiceDescriptor{Url: "testControlEndpointUrl"},
	}
	*loggingEndpoint = ""
	*artifactEndpoint = ""
	*controlEndpoint = ""
	err := ensureEndpointsSet(provisionInfo)
	if err != nil {
		t.Fatalf("ensureEndpointsSet() = %q, want nil", err)
	}
	if got, want := *loggingEndpoint, "testLoggingEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *loggingEndpoint = %q, want %q", got, want)
	}
	if got, want := *artifactEndpoint, "testArtifactEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *artifactEndpoint = %q, want %q", got, want)
	}
	if got, want := *controlEndpoint, "testControlEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *controlEndpoint = %q, want %q", got, want)
	}
}

func TestEnsureEndpointsSet_OneMissing(t *testing.T) {
	provisionInfo := &fnpb.ProvisionInfo{
		LoggingEndpoint:  &pipepb.ApiServiceDescriptor{Url: "testLoggingEndpointUrl"},
		ArtifactEndpoint: &pipepb.ApiServiceDescriptor{Url: "testArtifactEndpointUrl"},
		ControlEndpoint:  &pipepb.ApiServiceDescriptor{Url: ""},
	}
	*loggingEndpoint = ""
	*artifactEndpoint = ""
	*controlEndpoint = ""
	err := ensureEndpointsSet(provisionInfo)
	if err == nil {
		t.Fatalf("ensureEndpointsSet() = nil, want non-nil error")
	}
	if got, want := *loggingEndpoint, "testLoggingEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *loggingEndpoint = %q, want %q", got, want)
	}
	if got, want := *artifactEndpoint, "testArtifactEndpointUrl"; got != want {
		t.Fatalf("After ensureEndpointsSet(), *artifactEndpoint = %q, want %q", got, want)
	}
	if got, want := *controlEndpoint, ""; got != want {
		t.Fatalf("After ensureEndpointsSet(), *controlEndpoint = %q, want %q", got, want)
	}
}

func TestGetGoWorkerArtifactName_NoArtifacts(t *testing.T) {
	_, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, []*pipepb.ArtifactInformation{})
	if err == nil {
		t.Fatalf("getGoWorkerArtifactName() = nil, want non-nil error")
	}
}

func TestGetGoWorkerArtifactName_OneArtifact(t *testing.T) {
	artifact := constructArtifactInformation(t, artifact.URNGoWorkerBinaryRole, "test/path", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "test/path"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsFirstIsWorker(t *testing.T) {
	artifact1 := constructArtifactInformation(t, artifact.URNGoWorkerBinaryRole, "test/path", "sha")
	artifact2 := constructArtifactInformation(t, "other role", "test/path2", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "test/path"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsSecondIsWorker(t *testing.T) {
	artifact1 := constructArtifactInformation(t, "other role", "test/path", "sha")
	artifact2 := constructArtifactInformation(t, artifact.URNGoWorkerBinaryRole, "test/path2", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "test/path2"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsLegacyWay(t *testing.T) {
	artifact1 := constructArtifactInformation(t, "other role", "test/path", "sha")
	artifact2 := constructArtifactInformation(t, "other role", "worker", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	val, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err != nil {
		t.Fatalf("getGoWorkerArtifactName() = %v, want nil", err)
	}
	if got, want := val, "worker"; got != want {
		t.Fatalf("getGoWorkerArtifactName() = %v, want %v", got, want)
	}
}

func TestGetGoWorkerArtifactName_MultipleArtifactsNoneMatch(t *testing.T) {
	artifact1 := constructArtifactInformation(t, "other role", "test/path", "sha")
	artifact2 := constructArtifactInformation(t, "other role", "test/path2", "sha")
	artifacts := []*pipepb.ArtifactInformation{&artifact1, &artifact2}

	_, err := getGoWorkerArtifactName(context.Background(), &tools.Logger{}, artifacts)
	if err == nil {
		t.Fatalf("getGoWorkerArtifactName() = nil, want non-nil error")
	}
}

func TestCopyExe(t *testing.T) {
	testExeContent := []byte("testContent")

	// Make temp directory to cleanup at the end
	d, err := os.MkdirTemp(os.Getenv("TEST_TMPDIR"), "copyExe-*")
	if err != nil {
		t.Fatalf("failed to make temp directory, got %v", err)
	}
	t.Cleanup(func() { os.RemoveAll(d) })

	// Make our source file and write to it
	src, err := os.CreateTemp(d, "src.exe")
	if err != nil {
		t.Fatalf("failed to make temp file, got %v", err)
	}
	if _, err := src.Write(testExeContent); err != nil {
		t.Fatalf("failed to write to temp file, got %v", err)
	}
	if err := src.Close(); err != nil {
		t.Fatalf("failed to close temp file, got %v", err)
	}

	// Make sure our destination path doesn't exist already
	srcPath, destPath := src.Name(), filepath.Join(d, "dest.exe")
	if _, err := os.Stat(destPath); err == nil {
		t.Fatalf("dest file %v already exists", destPath)
	}

	err = copyExe(srcPath, destPath)
	if err != nil {
		t.Fatalf("copyExe() = %v, want nil", err)
	}
	if _, err := os.Stat(destPath); err != nil {
		t.Fatalf("After running copyExe, os.Stat() = %v, want nil", err)
	}
	destContents, err := os.ReadFile(destPath)
	if err != nil {
		t.Fatalf("After running copyExe, os.ReadFile() = %v, want nil", err)
	}
	if got, want := string(destContents), string(testExeContent); got != want {
		t.Fatalf("After running copyExe, os.ReadFile() = %v, want %v", got, want)
	}
}

func constructArtifactInformation(t *testing.T, roleUrn string, path string, sha string) pipepb.ArtifactInformation {
	t.Helper()

	typePayload, _ := proto.Marshal(&pipepb.ArtifactFilePayload{Path: path, Sha256: sha})

	return pipepb.ArtifactInformation{
		RoleUrn:     roleUrn,
		TypeUrn:     artifact.URNFileArtifact,
		TypePayload: typePayload,
	}
}

func TestConfigureGoogleCloudProfilerEnvVars(t *testing.T) {
	tests := []struct {
		name          string
		inputMetadata map[string]string
		expectedName  string
		expectedID    string
		expectedError string
	}{
		{
			"nil metadata",
			nil,
			"",
			"",
			"enable_google_cloud_profiler is set to true, but no metadata is received from provision server, profiling will not be enabled",
		},
		{
			"missing name",
			map[string]string{"job_id": "12345"},
			"",
			"",
			"required job_name missing from metadata, profiling will not be enabled without it",
		},
		{
			"missing id",
			map[string]string{"job_name": "my_job"},
			"",
			"",
			"required job_id missing from metadata, profiling will not be enabled without it",
		},
		{
			"correct",
			map[string]string{"job_name": "my_job", "job_id": "42"},
			"my_job",
			"42",
			"",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			t.Cleanup(os.Clearenv)
			err := configureGoogleCloudProfilerEnvVars(context.Background(), &tools.Logger{}, test.inputMetadata)
			if err != nil {
				if got, want := err.Error(), test.expectedError; got != want {
					t.Errorf("got error %v, want error %v", got, want)
				}
			}
			if got, want := os.Getenv(cloudProfilingJobName), test.expectedName; got != want {
				t.Errorf("got job name %v, want %v", got, want)
			}
			if got, want := os.Getenv(cloudProfilingJobID), test.expectedID; got != want {
				t.Errorf("got job id %v, want %v", got, want)
			}
		})
	}
}