Repository: G-Research/fasttrackml
Pull request #583 (2023-11-13 15:37:35)

Namespace integration for get artifacts endpoint
This PR is the third in a series to address #511 

What's inside the box:

- connect `namespace` logic to `GET /artifacts/get` endpoint.
- add missing unit tests.
- add missing `flow` integration tests.


Patch:
diff --git a/pkg/api/mlflow/controller/artifact.go b/pkg/api/mlflow/controller/artifact.go
index 63961a4..9fc7e17 100644
--- a/pkg/api/mlflow/controller/artifact.go
+++ b/pkg/api/mlflow/controller/artifact.go
@@ -50,7 +50,13 @@ func (c Controller) GetArtifact(ctx *fiber.Ctx) error {
 	}
 	log.Debugf("GetArtifact request: %#v", req)
 
-	artifact, err := c.artifactService.GetArtifact(ctx.Context(), &req)
+	ns, err := namespace.GetNamespaceFromContext(ctx.Context())
+	if err != nil {
+		return api.NewInternalError("error getting namespace from context")
+	}
+	log.Debugf("getArtifact namespace: %s", ns.Code)
+
+	artifact, err := c.artifactService.GetArtifact(ctx.Context(), ns, &req)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/api/mlflow/service/artifact/service.go b/pkg/api/mlflow/service/artifact/service.go
index f032eb2..3184d90 100644
--- a/pkg/api/mlflow/service/artifact/service.go
+++ b/pkg/api/mlflow/service/artifact/service.go
@@ -70,13 +70,13 @@ func (s Service) ListArtifacts(
 
 // GetArtifact handles business logic of `GET /artifacts/get` endpoint.
 func (s Service) GetArtifact(
-	ctx context.Context, req *request.GetArtifactRequest,
+	ctx context.Context, namespace *models.Namespace, req *request.GetArtifactRequest,
 ) (io.ReadCloser, error) {
 	if err := ValidateGetArtifactRequest(req); err != nil {
 		return nil, err
 	}
 
-	run, err := s.runRepository.GetByID(ctx, req.GetRunID())
+	run, err := s.runRepository.GetByNamespaceIDAndRunID(ctx, namespace.ID, req.GetRunID())
 	if err != nil {
 		return nil, api.NewInternalError("unable to find run '%s': %s", req.GetRunID(), err)
 	}


Test patch:
diff --git a/pkg/api/mlflow/service/artifact/service_test.go b/pkg/api/mlflow/service/artifact/service_test.go
index 2f5b1e0..063995f 100644
--- a/pkg/api/mlflow/service/artifact/service_test.go
+++ b/pkg/api/mlflow/service/artifact/service_test.go
@@ -1,8 +1,11 @@
 package artifact
 
 import (
+	"bytes"
 	"context"
 	"errors"
+	"io"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -178,3 +181,182 @@ func TestService_ListArtifacts_Error(t *testing.T) {
 		})
 	}
 }
+
+func TestService_GetArtifact_Ok(t *testing.T) {
+	artifactStorage := storage.MockArtifactStorageProvider{}
+	artifactStorage.On(
+		"Get", context.TODO(), "/artifact/uri", "",
+	).Return(
+		io.NopCloser(strings.NewReader("content")), nil,
+	)
+
+	artifactStorageFactory := storage.MockArtifactStorageFactoryProvider{}
+	artifactStorageFactory.On(
+		"GetStorage", context.TODO(), "/artifact/uri",
+	).Return(&artifactStorage, nil)
+
+	// init repository mocks.
+	runRepository := repositories.MockRunRepositoryProvider{}
+	runRepository.On(
+		"GetByNamespaceIDAndRunID",
+		context.TODO(),
+		uint(1),
+		"id",
+	).Return(&models.Run{
+		ID:          "id",
+		ArtifactURI: "/artifact/uri",
+	}, nil)
+
+	// call service under testing.
+	service := NewService(&runRepository, &artifactStorageFactory)
+	data, err := service.GetArtifact(
+		context.TODO(),
+		&models.Namespace{
+			ID: 1,
+		},
+		&request.GetArtifactRequest{
+			RunID: "id",
+		},
+	)
+
+	assert.Nil(t, err)
+	result := new(bytes.Buffer)
+	_, err = result.ReadFrom(data)
+	assert.Nil(t, err)
+	assert.Equal(t, "content", result.String())
+}
+
+func TestService_GetArtifact_Error(t *testing.T) {
+	testData := []struct {
+		name    string
+		error   *api.ErrorResponse
+		request *request.GetArtifactRequest
+		service func() *Service
+	}{
+		{
+			name:    "EmptyOrIncorrectRunID",
+			error:   api.NewInvalidParameterValueError("Missing value for required parameter 'run_id'"),
+			request: &request.GetArtifactRequest{},
+			service: func() *Service {
+				return NewService(
+					&repositories.MockRunRepositoryProvider{},
+					&storage.MockArtifactStorageFactoryProvider{},
+				)
+			},
+		},
+		{
+			name:  "PathIsRelativeAndContains2Dots",
+			error: api.NewInvalidParameterValueError("provided 'path' parameter is invalid"),
+			request: &request.GetArtifactRequest{
+				RunID: "id",
+				Path:  "../",
+			},
+			service: func() *Service {
+				return NewService(
+					&repositories.MockRunRepositoryProvider{},
+					&storage.MockArtifactStorageFactoryProvider{},
+				)
+			},
+		},
+		{
+			name:  "RunNotFoundDatabaseError",
+			error: api.NewInternalError("unable to find run 'id': database error"),
+			request: &request.GetArtifactRequest{
+				RunID: "id",
+			},
+			service: func() *Service {
+				runRepository := repositories.MockRunRepositoryProvider{}
+				runRepository.On(
+					"GetByNamespaceIDAndRunID",
+					context.TODO(),
+					uint(1),
+					"id",
+				).Return(nil, errors.New("database error"))
+				return NewService(
+					&runRepository,
+					&storage.MockArtifactStorageFactoryProvider{},
+				)
+			},
+		},
+		{
+			name:  "StorageError",
+			error: api.NewInternalError("error getting artifact object for URI: /artifact/uri"),
+			request: &request.GetArtifactRequest{
+				RunID: "id",
+			},
+			service: func() *Service {
+				artifactStorage := storage.MockArtifactStorageProvider{}
+				artifactStorage.On(
+					"Get", context.TODO(), "/artifact/uri", "",
+				).Return(
+					nil, errors.New("storage error"),
+				)
+
+				artifactStorageFactory := storage.MockArtifactStorageFactoryProvider{}
+				artifactStorageFactory.On(
+					"GetStorage", context.TODO(), "/artifact/uri",
+				).Return(&artifactStorage, nil)
+
+				runRepository := repositories.MockRunRepositoryProvider{}
+				runRepository.On(
+					"GetByNamespaceIDAndRunID",
+					context.TODO(),
+					uint(1),
+					"id",
+				).Return(&models.Run{
+					ID:          "id",
+					ArtifactURI: "/artifact/uri",
+				}, nil)
+				return NewService(
+					&runRepository,
+					&artifactStorageFactory,
+				)
+			},
+		},
+		{
+			name:  "UnsupportedStorage",
+			error: api.NewInternalError("run with id 'id' has unsupported artifact storage"),
+			request: &request.GetArtifactRequest{
+				RunID: "id",
+			},
+			service: func() *Service {
+				artifactStorage := storage.MockArtifactStorageProvider{}
+				artifactStorage.On(
+					"Get", context.TODO(), "/artifact/uri", "",
+				).Return(
+					nil, errors.New("storage error"),
+				)
+
+				artifactStorageFactory := storage.MockArtifactStorageFactoryProvider{}
+				artifactStorageFactory.On(
+					"GetStorage", context.TODO(), "/artifact/uri",
+				).Return(nil, errors.New("unsupported error"))
+
+				runRepository := repositories.MockRunRepositoryProvider{}
+				runRepository.On(
+					"GetByNamespaceIDAndRunID",
+					context.TODO(),
+					uint(1),
+					"id",
+				).Return(&models.Run{
+					ID:          "id",
+					ArtifactURI: "/artifact/uri",
+				}, nil)
+				return NewService(
+					&runRepository,
+					&artifactStorageFactory,
+				)
+			},
+		},
+	}
+
+	for _, tt := range testData {
+		t.Run(tt.name, func(t *testing.T) {
+			// call service under testing.
+			_, err := tt.service().GetArtifact(context.TODO(), &models.Namespace{
+				ID: 1,
+			}, tt.request)
+			assert.Equal(t, tt.error, err)
+		})
+	}
+}
diff --git a/tests/integration/golang/mlflow/namespace/flows/artifact_test.go b/tests/integration/golang/mlflow/namespace/flows/artifact_test.go
index b7b85a4..33cb1e9 100644
--- a/tests/integration/golang/mlflow/namespace/flows/artifact_test.go
+++ b/tests/integration/golang/mlflow/namespace/flows/artifact_test.go
@@ -3,6 +3,7 @@
 package flows
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"net/http"
@@ -233,6 +234,64 @@ func (s *ArtifactFlowTestSuite) testRunArtifactFlow(
 	)
 	assert.Equal(s.T(), fmt.Sprintf("RESOURCE_DOES_NOT_EXIST: unable to find run '%s'", run2ID), resp.Error())
 	assert.Equal(s.T(), api.ErrorCodeResourceDoesNotExist, string(resp.ErrorCode))
+
+	// test `GET /artifacts/get` endpoint.
+	s.getRunArtifactAndCompare(namespace1Code, request.GetArtifactRequest{
+		RunID: run1ID,
+		Path:  "artifact1.file",
+	}, "content1")
+
+	s.getRunArtifactAndCompare(namespace2Code, request.GetArtifactRequest{
+		RunID: run2ID,
+		Path:  "artifact2.file",
+	}, "content2")
+
+	// test `GET /artifacts/get` endpoint.
+	// check that there is no intersection between runs, so when we request
+	// run 1 in scope of namespace 2 and run 2 in scope of namespace 1 API will throw an error.
+	resp = api.ErrorResponse{}
+	assert.Nil(
+		s.T(),
+		s.MlflowClient.WithMethod(
+			http.MethodGet,
+		).WithNamespace(
+			namespace2Code,
+		).WithQuery(
+			request.GetArtifactRequest{
+				RunID: run1ID,
+			},
+		).WithResponseType(
+			helpers.ResponseTypeJSON,
+		).WithResponse(
+			&resp,
+		).DoRequest(
+			fmt.Sprintf("%s%s", mlflow.ArtifactsRoutePrefix, mlflow.ArtifactsGetRoute),
+		),
+	)
+	assert.Equal(s.T(), fmt.Sprintf("RESOURCE_DOES_NOT_EXIST: unable to find run '%s'", run1ID), resp.Error())
+	assert.Equal(s.T(), api.ErrorCodeResourceDoesNotExist, string(resp.ErrorCode))
+
+	resp = api.ErrorResponse{}
+	assert.Nil(
+		s.T(),
+		s.MlflowClient.WithMethod(
+			http.MethodGet,
+		).WithNamespace(
+			namespace1Code,
+		).WithQuery(
+			request.ListArtifactsRequest{
+				RunID: run2ID,
+			},
+		).WithResponseType(
+			helpers.ResponseTypeJSON,
+		).WithResponse(
+			&resp,
+		).DoRequest(
+			fmt.Sprintf("%s%s", mlflow.ArtifactsRoutePrefix, mlflow.ArtifactsGetRoute),
+		),
+	)
+	assert.Equal(s.T(), fmt.Sprintf("RESOURCE_DOES_NOT_EXIST: unable to find run '%s'", run2ID), resp.Error())
+	assert.Equal(s.T(), api.ErrorCodeResourceDoesNotExist, string(resp.ErrorCode))
 }
 
 func (s *ArtifactFlowTestSuite) createRun(namespace string, req *request.CreateRunRequest) string {
@@ -278,3 +337,23 @@ func (s *ArtifactFlowTestSuite) listRunArtifactsAndCompare(
 	)
 	assert.Equal(s.T(), expectedResponse, actualResponse.Files)
 }
+
+func (s *ArtifactFlowTestSuite) getRunArtifactAndCompare(
+	namespace string, req request.GetArtifactRequest, expectedResponse string,
+) {
+	actualResponse := new(bytes.Buffer)
+	assert.Nil(s.T(), s.MlflowClient.WithMethod(
+		http.MethodGet,
+	).WithNamespace(
+		namespace,
+	).WithQuery(
+		req,
+	).WithResponseType(
+		helpers.ResponseTypeBuffer,
+	).WithResponse(
+		actualResponse,
+	).DoRequest(
+		fmt.Sprintf("%s%s", mlflow.ArtifactsRoutePrefix, mlflow.ArtifactsGetRoute),
+	))
+	assert.Equal(s.T(), expectedResponse, actualResponse.String())
+}