feat(master): record image versions and track poll suggestions in the store (#9)

internal/e2e/agent_master_test.go

@@ -44,6 +44,51 @@ 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
 }
 
 // seedService adds a service row that jobs can be triggered for.
@@ -88,7 +133,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
 }
@@ -185,9 +230,10 @@ func (m *memStore) FinishJob(_ context.Context, agentID, jobID int64, f model.Jo
 		svc := m.svcRows[j.ServiceID]
 		svc.DesiredRef = f.Ref
 		// 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
 	}