feat: master/agent phase 4 — registry polling and version history (#9) (#23)

internal/e2e/agent_master_test.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"regexp"
 	"sort"
 	"strings"
 	"sync"
@@ -31,6 +32,9 @@ import (
 	"github.com/shcizo/package-updater/internal/protocol"
 )
 
+// digestPattern is the digest format the real store records.
+var digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
+
 type memStore struct {
 	mu        sync.Mutex
 	agents    map[string]model.Agent
@@ -44,6 +48,63 @@ type memStore struct {
 	jobs          map[int64]*model.Job
 	nextJob       int64
 	finishApplied int // FinishJob calls that returned true
+	versions      map[int64]model.ImageVersion
+}
+
+// seedVersion adds an image_versions row that GetImageVersion returns.
+func (m *memStore) seedVersion(v model.ImageVersion) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.versions == nil {
+		m.versions = map[int64]model.ImageVersion{}
+	}
+	m.versions[v.ID] = v
+}
+
+func (m *memStore) GetImageVersion(_ context.Context, id int64) (model.ImageVersion, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	v, ok := m.versions[id]
+	if !ok {
+		return model.ImageVersion{}, model.ErrNotFound
+	}
+	return v, nil
+}
+
+func (m *memStore) ListServices(_ context.Context) ([]model.Service, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	var out []model.Service
+	for _, svc := range m.svcRows {
+		if svc.RemovedAt == nil {
+			out = append(out, svc)
+		}
+	}
+	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
+	return out, nil
+}
+
+func (m *memStore) SetPollError(_ context.Context, ids []int64, msg string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for _, id := range ids {
+		svc := m.svcRows[id]
+		svc.LastPollError = msg
+		m.svcRows[id] = svc
+	}
+	return nil
+}
+
+func (m *memStore) SetLastTagDigest(_ context.Context, id int64, digest string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	svc, ok := m.svcRows[id]
+	if !ok {
+		return model.ErrNotFound
+	}
+	svc.LastTagDigest = digest
+	m.svcRows[id] = svc
+	return nil
 }
 
 // seedService adds a service row that jobs can be triggered for.
@@ -88,7 +149,7 @@ func (m *memStore) SetAvailable(_ context.Context, id int64, a model.Available)
 	if !ok {
 		return model.ErrNotFound
 	}
-	svc.AvailableRef, svc.AvailableSource, svc.AvailableAt = a.Ref, a.Source, &a.At
+	svc.AvailableRef, svc.AvailableDigest, svc.AvailableSource, svc.AvailableAt = a.Ref, a.Digest, a.Source, &a.At
 	m.svcRows[id] = svc
 	return nil
 }
@@ -184,10 +245,14 @@ func (m *memStore) FinishJob(_ context.Context, agentID, jobID int64, f model.Jo
 	if f.Status == model.JobSucceeded {
 		svc := m.svcRows[j.ServiceID]
 		svc.DesiredRef = f.Ref
+		if _, d, _ := strings.Cut(f.Ref, "@"); !strings.Contains(requestedRef, "@") && digestPattern.MatchString(d) {
+			svc.LastTagDigest = d
+		}
 		// Same rule as the SQL: the job was created for the suggested ref after it was reported.
-		requested, _, _ := strings.Cut(requestedRef, "@")
-		if svc.AvailableAt != nil && svc.AvailableRef == requested && !svc.AvailableAt.After(j.CreatedAt) {
-			svc.AvailableRef, svc.AvailableSource, svc.AvailableAt = "", "", nil
+		requested, digest, _ := strings.Cut(requestedRef, "@")
+		if svc.AvailableAt != nil && svc.AvailableRef == requested && !svc.AvailableAt.After(j.CreatedAt) &&
+			(digest == "" || svc.AvailableDigest == digest) {
+			svc.AvailableRef, svc.AvailableDigest, svc.AvailableSource, svc.AvailableAt = "", "", "", nil
 		}
 		m.svcRows[j.ServiceID] = svc
 	}
@@ -282,7 +347,9 @@ type currentHub struct{ p atomic.Pointer[hub.Hub] }
 
 func (c *currentHub) ServeHTTP(w http.ResponseWriter, r *http.Request) { c.p.Load().ServeHTTP(w, r) }
 func (c *currentHub) Disconnect(id int64)                              { c.p.Load().Disconnect(id) }
-func (c *currentHub) Route(svc model.Service) (model.Route, error)     { return c.p.Load().Route(svc) }
+func (c *currentHub) Route(svc model.Service, minProtocol int) (model.Route, error) {
+	return c.p.Load().Route(svc, minProtocol)
+}
 func (c *currentHub) StartJob(r model.Route, id int64, s protocol.JobStart) error {
 	return c.p.Load().StartJob(r, id, s)
 }