package artifactregistryservice_test import ( "context" "fmt" "reflect" "testing" artifactRegistryService "github.com/BishopFox/cloudfox/gcp/services/artifactRegistryService" "github.com/BishopFox/cloudfox/gcp/services/models" artifactregistrypb "cloud.google.com/go/artifactregistry/apiv1/artifactregistrypb" "github.com/googleapis/gax-go/v2" "google.golang.org/api/iterator" locationpb "google.golang.org/genproto/googleapis/cloud/location" ) // Mock structure for ArtifactRegistryClientWrapper type mockArtifactRegistryClient struct { repositories []*artifactregistrypb.Repository locations []*locationpb.Location dockerImages []*artifactregistrypb.DockerImage } func (m *mockArtifactRegistryClient) Close() error { return nil } // Mock iterator for repositories type mockRepoIterator struct { repositories []*artifactregistrypb.Repository index int } func (it *mockRepoIterator) Next() (*artifactregistrypb.Repository, error) { if it.index < len(it.repositories) { repo := it.repositories[it.index] it.index++ return repo, nil } return nil, iterator.Done } // Mock iterator for locations type mockLocationIterator struct { locations []*locationpb.Location index int } func (it *mockLocationIterator) Next() (*locationpb.Location, error) { if it.index < len(it.locations) { loc := it.locations[it.index] it.index++ return loc, nil } return nil, iterator.Done } // Mock iterator for Docker images type mockDockerImageIterator struct { dockerImages []*artifactregistrypb.DockerImage index int } func (it *mockDockerImageIterator) Next() (*artifactregistrypb.DockerImage, error) { if it.index < len(it.dockerImages) { img := it.dockerImages[it.index] it.index++ return img, nil } return nil, iterator.Done } // Mock methods func (m *mockArtifactRegistryClient) ListRepositories(ctx context.Context, req *artifactregistrypb.ListRepositoriesRequest, opts ...gax.CallOption) models.GenericIterator[artifactregistrypb.Repository] { return &mockRepoIterator{repositories: m.repositories} } func (m *mockArtifactRegistryClient) ListLocations(ctx context.Context, req *locationpb.ListLocationsRequest, opts ...gax.CallOption) models.GenericIterator[locationpb.Location] { return &mockLocationIterator{locations: m.locations} } func (m *mockArtifactRegistryClient) GetRepository(ctx context.Context, req *artifactregistrypb.GetRepositoryRequest, opts ...gax.CallOption) (*artifactregistrypb.Repository, error) { for _, repo := range m.repositories { if repo.Name == req.Name { return repo, nil } } return nil, fmt.Errorf("repository not found") } func (m *mockArtifactRegistryClient) ListDockerImages(ctx context.Context, req *artifactregistrypb.ListDockerImagesRequest, opts ...gax.CallOption) models.GenericIterator[artifactregistrypb.DockerImage] { return &mockDockerImageIterator{dockerImages: m.dockerImages} } var mockClient = mockArtifactRegistryClient{} var ars = artifactRegistryService.ArtifactRegistryService{ Client: &artifactRegistryService.ArtifactRegistryClientWrapper{ Closer: mockClient.Close, RepositoryLister: func(ctx context.Context, req *artifactregistrypb.ListRepositoriesRequest, opts ...gax.CallOption) models.GenericIterator[artifactregistrypb.Repository] { return mockClient.ListRepositories(ctx, req, opts...) }, LocationLister: func(ctx context.Context, req *locationpb.ListLocationsRequest, opts ...gax.CallOption) models.GenericIterator[locationpb.Location] { return mockClient.ListLocations(ctx, req, opts...) }, RepositoryGetter: func(ctx context.Context, req *artifactregistrypb.GetRepositoryRequest, opts ...gax.CallOption) (*artifactregistrypb.Repository, error) { return mockClient.GetRepository(ctx, req, opts...) }, DockerImageLister: func(ctx context.Context, req *artifactregistrypb.ListDockerImagesRequest, opts ...gax.CallOption) models.GenericIterator[artifactregistrypb.DockerImage] { return mockClient.ListDockerImages(ctx, req, opts...) }, }, } // Tests go here func TestArtifacts(t *testing.T) { cases := []struct { name string projectID string location string repositoryName string setupMockClient func(client *mockArtifactRegistryClient) expectedArtifacts []artifactRegistryService.ArtifactInfo expectError bool }{ { name: "Valid artifacts retrieval", projectID: "project1", location: "us-central1", repositoryName: "repo1", setupMockClient: func(client *mockArtifactRegistryClient) { client.dockerImages = []*artifactregistrypb.DockerImage{ { Name: "projects/project1/locations/us-central1/repositories/repo1/dockerImages/image1@sha256:e9954c1fc875017be1c3e36eca16be2d9e9bccc4bf072163515467d6a823c7cf", Uri: "us-central1-docker.pkg.dev/project1/repo1/image1@sha256:e9954c1fc875017be1c3e36eca16be2d9e9bccc4bf072163515467d6a823c7cf", ImageSizeBytes: 1024, }, } client.locations = []*locationpb.Location{ { LocationId: "us-central1", }, } client.repositories = []*artifactregistrypb.Repository{{ Name: "projects/project1/locations/us-central1/repositories/repo1", Format: artifactregistrypb.Repository_DOCKER, }, } }, expectedArtifacts: []artifactRegistryService.ArtifactInfo{ { Name: "image1", Format: "DOCKER", Location: "us-central1", Repository: "repo1", SizeBytes: "1024", ProjectID: "project1", Digest: "sha256:e9954c1fc875017be1c3e36eca16be2d9e9bccc4bf072163515467d6a823c7cf", Updated: "1970-01-01 00:00:00 +0000 UTC", }, }, expectError: false, }, // Add more test cases here } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { // Reset mockClient state mockClient.repositories = nil mockClient.locations = nil mockClient.dockerImages = nil // Setup mock client for the current test case tc.setupMockClient(&mockClient) artifacts, err := ars.Artifacts(tc.projectID, tc.location, tc.repositoryName) if (err != nil) != tc.expectError { t.Errorf("Artifacts() error = %v, expectError %v", err, tc.expectError) return } if !reflect.DeepEqual(artifacts, tc.expectedArtifacts) { t.Errorf("Artifacts() = %v, expected %v", artifacts, tc.expectedArtifacts) } }) } } func TestRepositories(t *testing.T) { cases := []struct { name string projectID string setupMockClient func(client *mockArtifactRegistryClient) expectedRepositories []artifactRegistryService.RepositoryInfo expectError bool }{ { name: "Valid repositories retrieval", projectID: "project1", setupMockClient: func(client *mockArtifactRegistryClient) { client.repositories = []*artifactregistrypb.Repository{ { Name: "projects/project1/locations/us-central1/repositories/repo1", Format: artifactregistrypb.Repository_DOCKER, Description: "Test repository", }, } client.locations = []*locationpb.Location{ { LocationId: "us-central1", }, } }, expectedRepositories: []artifactRegistryService.RepositoryInfo{ { Name: "projects/project1/locations/us-central1/repositories/repo1", Format: "DOCKER", Description: "Test repository", SizeBytes: "0", ProjectID: "project1", Location: "us-central1", }, }, expectError: false, }, // Add more test cases here } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { // Reset mockClient state mockClient.repositories = nil mockClient.locations = nil mockClient.dockerImages = nil // Setup mock client for the current test case tc.setupMockClient(&mockClient) repositories, err := ars.Repositories(tc.projectID) if (err != nil) != tc.expectError { t.Errorf("Repositories() error = %v, expectError %v", err, tc.expectError) return } if !reflect.DeepEqual(repositories, tc.expectedRepositories) { t.Errorf("Repositories() = %v, expected %v", repositories, tc.expectedRepositories) } }) } }