diff --git a/go.mod b/go.mod index 3ef9ba8..9e04148 100644 --- a/go.mod +++ b/go.mod @@ -19,6 +19,7 @@ require ( github.com/urfave/cli/v3 v3.11.0 golang.org/x/sync v0.22.0 golang.org/x/sys v0.47.0 + google.golang.org/grpc v1.83.2 gopkg.in/yaml.v3 v3.0.1 ) @@ -82,7 +83,6 @@ require ( golang.org/x/net v0.58.0 // indirect golang.org/x/text v0.41.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260904194346-d0f1323225a4 // indirect - google.golang.org/grpc v1.83.2 // indirect google.golang.org/protobuf v1.36.12 // indirect gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect ) diff --git a/store/core/workload.go b/store/core/workload.go index e0ed2ee..21b9a83 100644 --- a/store/core/workload.go +++ b/store/core/workload.go @@ -10,7 +10,8 @@ import ( "time" pb "github.com/projecteru2/core/rpc/gen" - coretypes "github.com/projecteru2/core/types" + "google.golang.org/grpc/codes" + grpcstatus "google.golang.org/grpc/status" "github.com/projecteru2/agent/types" ) @@ -37,13 +38,7 @@ func (s *Store) WorkloadExists(ctx context.Context, ID string) (bool, error) { _, err := call(ctx, s, func(ctx context.Context) (*pb.Workload, error) { return s.client().GetWorkload(ctx, &pb.WorkloadID{Id: ID}) }) - switch { - case err == nil: - return true, nil - case strings.Contains(err.Error(), coretypes.ErrInvaildCount.Error()), strings.Contains(err.Error(), coretypes.ErrWorkloadNotExists.Error()): - return false, nil - } - return false, err + return workloadLookupResult(err) } func (s *Store) SetWorkloadStatus(ctx context.Context, status *types.WorkloadStatus) error { @@ -94,3 +89,13 @@ func getCacheTTL(ttl int64) time.Duration { delta := rand.Int64N(max(ttl, 1)) / 4 //nolint:gosec // cache ttl jitter needs no csprng return time.Duration(ttl-ttl/8+delta) * time.Second } + +func workloadLookupResult(err error) (bool, error) { + switch { + case err == nil: + return true, nil + case grpcstatus.Code(err) == codes.NotFound: + return false, nil + } + return false, err +} diff --git a/store/core/workload_test.go b/store/core/workload_test.go new file mode 100644 index 0000000..6d66a9c --- /dev/null +++ b/store/core/workload_test.go @@ -0,0 +1,33 @@ +package core + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + grpcstatus "google.golang.org/grpc/status" +) + +func TestWorkloadLookupResult(t *testing.T) { + transport := errors.New("dial core: connection refused") + legacy := grpcstatus.Error(1051, "workload not exists") + tests := []struct { + name string + err error + exists bool + wantErr error + }{ + {"found", nil, true, nil}, + {"not found", grpcstatus.Error(codes.NotFound, "workload not exists"), false, nil}, + {"per-rpc code", legacy, false, legacy}, + {"transport", transport, false, transport}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + exists, err := workloadLookupResult(tt.err) + assert.Equal(t, tt.exists, exists) + assert.Equal(t, tt.wantErr, err) + }) + } +}