Skip to content
Merged
Prev Previous commit
feat: add test-code for s3 registry store
Signed-off-by: iamcodingcat <joyh951021@gmail.com>
  • Loading branch information
iamcodingcat committed May 13, 2025
commit 486262263fc8e5e7be21fb3e908be3840badc622
94 changes: 94 additions & 0 deletions go/internal/feast/registry/registry_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
package registry

import (
"context"
"errors"
"io/ioutil"
"net/url"
"strings"
"testing"
"time"

"github.com/aws/aws-sdk-go-v2/service/s3"
)

func TestGetOnlineFeaturesS3Registry(t *testing.T) {
mockS3Client := &MockS3Client{
GetObjectFn: func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
return &s3.GetObjectOutput{
Body: ioutil.NopCloser(strings.NewReader("mock data")),
}, nil
},
DeleteObjectFn: func(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error) {
return &s3.DeleteObjectOutput{}, nil
},
}

tests := []struct {
name string
config *RepoConfig
}{
{
name: "redis with simple features",
config: &RepoConfig{
Project: "feature_repo",
Registry: map[string]interface{}{
"path": "s3://test-bucket/path/to/registry.db",
},
Provider: "aws",
},
},
}
for _, test := range tests {
registryConfig, err := test.config.GetRegistryConfig()
if err != nil {
t.Errorf("Error getting registry config. msg: %s", err.Error())
}
r := &Registry{
project: test.config.Project,
cachedRegistryProtoTtl: time.Duration(registryConfig.CacheTtlSeconds) * time.Second,
}
_ = registryConfig.RegistryStoreType
registryPath := registryConfig.Path
uri, err := url.Parse(registryPath)
if err != nil {
t.Errorf("Error parsing registry path. msg: %s", err.Error())
}
if registryStoreType, ok := REGISTRY_STORE_CLASS_FOR_SCHEME[uri.Scheme]; ok {
switch registryStoreType {
case "S3RegistryStore":
registryStore := &S3RegistryStore{
filePath: registryConfig.Path,
s3Client: mockS3Client,
}
r.registryStore = registryStore
err := r.InitializeRegistry()
if err != nil {
t.Errorf("Error initializing registry. msg: %s. registry path=%q", err.Error(), registryPath)
}
default:
t.Errorf("Only S3RegistryStore is supported on this testing. got=%s", registryStoreType)
}
}
}
}

// MockS3Client is mock client for testing s3 registry store
type MockS3Client struct {
GetObjectFn func(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
DeleteObjectFn func(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error)
}

func (m *MockS3Client) GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
if m.GetObjectFn != nil {
return m.GetObjectFn(ctx, params)
}
return nil, errors.New("not implemented")
}

func (m *MockS3Client) DeleteObject(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error) {
if m.DeleteObjectFn != nil {
return m.DeleteObjectFn(ctx, params)
}
return nil, errors.New("not implemented")
}
12 changes: 9 additions & 3 deletions go/internal/feast/registry/s3.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,16 @@ import (
"google.golang.org/protobuf/proto"
)

// S3ClientInterface define interface of s3.Client for making mocking s3 client and testing it
type S3ClientInterface interface {
GetObject(ctx context.Context, params *s3.GetObjectInput, optFns ...func(*s3.Options)) (*s3.GetObjectOutput, error)
DeleteObject(ctx context.Context, params *s3.DeleteObjectInput, optFns ...func(*s3.Options)) (*s3.DeleteObjectOutput, error)
}

// A S3RegistryStore is a S3 object storage-based implementation of the RegistryStore interface
type S3RegistryStore struct {
filePath string
s3Client *s3.Client
s3Client S3ClientInterface
}

// NewS3RegistryStore creates a S3RegistryStore with the given configuration
Expand Down Expand Up @@ -55,7 +61,7 @@ func (r *S3RegistryStore) GetRegistryProto() (*core.Registry, error) {
Key: aws.String(key),
})
if err != nil {
return nil, err
panic(err)
}
defer output.Body.Close()

Expand Down Expand Up @@ -88,7 +94,7 @@ func (r *S3RegistryStore) Teardown() error {
Key: aws.String(key),
})
if err != nil {
return err
panic(err)
}
return nil
}
Expand Down
Loading