mirror of
https://github.com/coder/coder.git
synced 2026-09-21 20:51:01 +08:00
feat(agent): add agent unit manager (#20715)
relates to: https://github.com/coder/internal/issues/1094 This is number 1 of 5 pull requests in an effort to add agent script ordering. It adds a unit manager, which uses an underlying DAG and a list of subscribers to inform units when their dependencies have changed in status. In follow-up PRs: * This unit manager will be plumbed into the workspace agent struct. * It will then be exposed to users via a new socket based drpc API * The agentsocket API will then become accessible via CLI commands that allow coder scripts to express their dependencies on one another. This is an experimental feature. There may be ways to improve the efficiency of the manager struct, but it is more important to validate this feature with customers before we invest in such optimizations. See the tests for examples of how units may communicate with one another. Actual CLI usage will be analogous. I used an LLM to produce some of these changes, but I have conducted thorough self review and consider this contribution to be ready for an external reviewer.
This commit is contained in:
Sas Swart
+1
-1
@@ -58,7 +58,7 @@ func (g *Graph[EdgeType, VertexType]) AddEdge(from, to VertexType, edge EdgeType
|
||||
toID := g.getOrCreateVertexID(to)
|
||||
|
||||
if g.canReach(to, from) {
|
||||
return xerrors.Errorf("adding edge (%v -> %v) would create a cycle", from, to)
|
||||
return xerrors.Errorf("adding edge (%v -> %v): %w", from, to, ErrCycleDetected)
|
||||
}
|
||||
|
||||
g.gonumGraph.SetEdge(simple.Edge{F: simple.Node(fromID), T: simple.Node(toID)})
|
||||
|
||||
@@ -148,8 +148,7 @@ func TestGraph(t *testing.T) {
|
||||
graph := &testGraph{}
|
||||
unit1 := &testGraphVertex{Name: "unit1"}
|
||||
err := graph.AddEdge(unit1, unit1, testEdgeCompleted)
|
||||
require.Error(t, err)
|
||||
require.ErrorContains(t, err, fmt.Sprintf("adding edge (%v -> %v) would create a cycle", unit1, unit1))
|
||||
require.ErrorIs(t, err, unit.ErrCycleDetected)
|
||||
|
||||
return graph
|
||||
},
|
||||
@@ -160,8 +159,7 @@ func TestGraph(t *testing.T) {
|
||||
err := graph.AddEdge(unit1, unit2, testEdgeCompleted)
|
||||
require.NoError(t, err)
|
||||
err = graph.AddEdge(unit2, unit1, testEdgeStarted)
|
||||
require.Error(t, err)
|
||||
require.ErrorContains(t, err, fmt.Sprintf("adding edge (%v -> %v) would create a cycle", unit2, unit1))
|
||||
require.ErrorIs(t, err, unit.ErrCycleDetected)
|
||||
|
||||
return graph
|
||||
},
|
||||
@@ -341,7 +339,7 @@ func TestGraphThreadSafety(t *testing.T) {
|
||||
// Verify all attempts correctly returned cycle error
|
||||
for i, err := range cycleErrors {
|
||||
require.Error(t, err, "goroutine %d should have detected cycle", i)
|
||||
require.Contains(t, err.Error(), "would create a cycle")
|
||||
require.ErrorIs(t, err, unit.ErrCycleDetected)
|
||||
}
|
||||
|
||||
// Verify graph remains valid (original chain intact)
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
package unit
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/util/slice"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrUnitNotFound = xerrors.New("unit not found")
|
||||
ErrUnitAlreadyRegistered = xerrors.New("unit already registered")
|
||||
ErrCannotUpdateOtherUnit = xerrors.New("cannot update other unit's status")
|
||||
ErrDependenciesNotSatisfied = xerrors.New("unit dependencies not satisfied")
|
||||
ErrSameStatusAlreadySet = xerrors.New("same status already set")
|
||||
ErrCycleDetected = xerrors.New("cycle detected")
|
||||
ErrFailedToAddDependency = xerrors.New("failed to add dependency")
|
||||
)
|
||||
|
||||
// Status represents the status of a unit.
|
||||
type Status string
|
||||
|
||||
// Status constants for dependency tracking.
|
||||
const (
|
||||
StatusNotRegistered Status = ""
|
||||
StatusPending Status = "pending"
|
||||
StatusStarted Status = "started"
|
||||
StatusComplete Status = "completed"
|
||||
)
|
||||
|
||||
// ID provides a type narrowed representation of the unique identifier of a unit.
|
||||
type ID string
|
||||
|
||||
// Unit represents a point-in-time snapshot of a vertex in the dependency graph.
|
||||
// Units may depend on other units, or be depended on by other units. The unit struct
|
||||
// is not aware of updates made to the dependency graph after it is initialized and should
|
||||
// not be cached.
|
||||
type Unit struct {
|
||||
id ID
|
||||
status Status
|
||||
// ready is true if all dependencies are satisfied.
|
||||
// It does not have an accessor method on Unit, because a unit cannot know whether it is ready.
|
||||
// Only the Manager can calculate whether a unit is ready based on knowledge of the dependency graph.
|
||||
// To discourage use of an outdated readiness value, only the Manager should set and return this field.
|
||||
ready bool
|
||||
}
|
||||
|
||||
func (u Unit) ID() ID {
|
||||
return u.id
|
||||
}
|
||||
|
||||
func (u Unit) Status() Status {
|
||||
return u.status
|
||||
}
|
||||
|
||||
// Dependency represents a dependency relationship between units.
|
||||
type Dependency struct {
|
||||
Unit ID
|
||||
DependsOn ID
|
||||
RequiredStatus Status
|
||||
CurrentStatus Status
|
||||
IsSatisfied bool
|
||||
}
|
||||
|
||||
// Manager provides reactive dependency tracking over a Graph.
|
||||
// It manages Unit registration, dependency relationships, and status updates
|
||||
// with automatic recalculation of readiness when dependencies are satisfied.
|
||||
type Manager struct {
|
||||
mu sync.RWMutex
|
||||
|
||||
// The underlying graph that stores dependency relationships
|
||||
graph *Graph[Status, ID]
|
||||
|
||||
// Store vertex instances for each unit to ensure consistent references
|
||||
units map[ID]Unit
|
||||
}
|
||||
|
||||
// NewManager creates a new Manager instance.
|
||||
func NewManager() *Manager {
|
||||
return &Manager{
|
||||
graph: &Graph[Status, ID]{},
|
||||
units: make(map[ID]Unit),
|
||||
}
|
||||
}
|
||||
|
||||
// Register adds a unit to the manager if it is not already registered.
|
||||
// If a Unit is already registered (per the ID field), it is not updated.
|
||||
func (m *Manager) Register(id ID) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if m.registered(id) {
|
||||
return xerrors.Errorf("registering unit %q: %w", id, ErrUnitAlreadyRegistered)
|
||||
}
|
||||
|
||||
m.units[id] = Unit{
|
||||
id: id,
|
||||
status: StatusPending,
|
||||
ready: true,
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// registered checks if a unit is registered in the manager.
|
||||
func (m *Manager) registered(id ID) bool {
|
||||
return m.units[id].status != StatusNotRegistered
|
||||
}
|
||||
|
||||
// Unit fetches a unit from the manager. If the unit does not exist,
|
||||
// it returns the Unit zero-value as a placeholder unit, because
|
||||
// units may depend on other units that have not yet been created.
|
||||
func (m *Manager) Unit(id ID) Unit {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
return m.units[id]
|
||||
}
|
||||
|
||||
func (m *Manager) IsReady(id ID) bool {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
if !m.registered(id) {
|
||||
return false
|
||||
}
|
||||
|
||||
return m.units[id].ready
|
||||
}
|
||||
|
||||
// AddDependency adds a dependency relationship between units.
|
||||
// The unit depends on the dependsOn unit reaching the requiredStatus.
|
||||
func (m *Manager) AddDependency(unit ID, dependsOn ID, requiredStatus Status) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if !m.registered(unit) {
|
||||
return xerrors.Errorf("checking registration for unit %q: %w", unit, ErrUnitNotFound)
|
||||
}
|
||||
|
||||
// Add the dependency edge to the graph
|
||||
// The edge goes from unit to dependsOn, representing the dependency
|
||||
err := m.graph.AddEdge(unit, dependsOn, requiredStatus)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("adding edge for unit %q: %w", unit, errors.Join(ErrFailedToAddDependency, err))
|
||||
}
|
||||
|
||||
// Recalculate readiness for the unit since it now has a new dependency
|
||||
m.recalculateReadinessUnsafe(unit)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateStatus updates a unit's status and recalculates readiness for affected dependents.
|
||||
func (m *Manager) UpdateStatus(unit ID, newStatus Status) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if !m.registered(unit) {
|
||||
return xerrors.Errorf("checking registration for unit %q: %w", unit, ErrUnitNotFound)
|
||||
}
|
||||
|
||||
u := m.units[unit]
|
||||
if u.status == newStatus {
|
||||
return xerrors.Errorf("checking status for unit %q: %w", unit, ErrSameStatusAlreadySet)
|
||||
}
|
||||
|
||||
u.status = newStatus
|
||||
m.units[unit] = u
|
||||
|
||||
// Get all units that depend on this one (reverse adjacent vertices)
|
||||
dependents := m.graph.GetReverseAdjacentVertices(unit)
|
||||
|
||||
// Recalculate readiness for all dependents
|
||||
for _, dependent := range dependents {
|
||||
m.recalculateReadinessUnsafe(dependent.From)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// recalculateReadinessUnsafe recalculates the readiness state for a unit.
|
||||
// This method assumes the caller holds the write lock.
|
||||
func (m *Manager) recalculateReadinessUnsafe(unit ID) {
|
||||
u := m.units[unit]
|
||||
dependencies := m.graph.GetForwardAdjacentVertices(unit)
|
||||
|
||||
allSatisfied := true
|
||||
for _, dependency := range dependencies {
|
||||
requiredStatus := dependency.Edge
|
||||
dependsOnUnit := m.units[dependency.To]
|
||||
if dependsOnUnit.status != requiredStatus {
|
||||
allSatisfied = false
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
u.ready = allSatisfied
|
||||
m.units[unit] = u
|
||||
}
|
||||
|
||||
// GetGraph returns the underlying graph for visualization and debugging.
|
||||
// This should be used carefully as it exposes the internal graph structure.
|
||||
func (m *Manager) GetGraph() *Graph[Status, ID] {
|
||||
return m.graph
|
||||
}
|
||||
|
||||
// GetAllDependencies returns all dependencies for a unit, both satisfied and unsatisfied.
|
||||
func (m *Manager) GetAllDependencies(unit ID) ([]Dependency, error) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
if !m.registered(unit) {
|
||||
return nil, xerrors.Errorf("checking registration for unit %q: %w", unit, ErrUnitNotFound)
|
||||
}
|
||||
|
||||
dependencies := m.graph.GetForwardAdjacentVertices(unit)
|
||||
|
||||
var allDependencies []Dependency
|
||||
|
||||
for _, dependency := range dependencies {
|
||||
dependsOnUnit := m.units[dependency.To]
|
||||
requiredStatus := dependency.Edge
|
||||
allDependencies = append(allDependencies, Dependency{
|
||||
Unit: unit,
|
||||
DependsOn: dependsOnUnit.id,
|
||||
RequiredStatus: requiredStatus,
|
||||
CurrentStatus: dependsOnUnit.status,
|
||||
IsSatisfied: dependsOnUnit.status == requiredStatus,
|
||||
})
|
||||
}
|
||||
|
||||
return allDependencies, nil
|
||||
}
|
||||
|
||||
// GetUnmetDependencies returns a list of unsatisfied dependencies for a unit.
|
||||
func (m *Manager) GetUnmetDependencies(unit ID) ([]Dependency, error) {
|
||||
allDependencies, err := m.GetAllDependencies(unit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var unmetDependencies []Dependency = slice.Filter(allDependencies, func(dependency Dependency) bool {
|
||||
return !dependency.IsSatisfied
|
||||
})
|
||||
|
||||
return unmetDependencies, nil
|
||||
}
|
||||
|
||||
// ExportDOT exports the dependency graph to DOT format for visualization.
|
||||
func (m *Manager) ExportDOT(name string) (string, error) {
|
||||
return m.graph.ToDOT(name)
|
||||
}
|
||||
@@ -0,0 +1,594 @@
|
||||
package unit_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/agent/unit"
|
||||
)
|
||||
|
||||
const (
|
||||
unitA unit.ID = "serviceA"
|
||||
unitB unit.ID = "serviceB"
|
||||
unitC unit.ID = "serviceC"
|
||||
unitD unit.ID = "serviceD"
|
||||
)
|
||||
|
||||
func TestManager_Register(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("RegisterNewUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: a unit is registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: the unit should be ready (no dependencies)
|
||||
u := manager.Unit(unitA)
|
||||
assert.Equal(t, unitA, u.ID())
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("RegisterDuplicateUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: a unit is registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Newly registered units have StatusPending. We update the unit status to StatusStarted,
|
||||
// so we can later assert that it is not overwritten back to StatusPending by the second
|
||||
// register call
|
||||
manager.UpdateStatus(unitA, unit.StatusStarted)
|
||||
|
||||
// When: the unit is registered again
|
||||
err = manager.Register(unitA)
|
||||
|
||||
// Then: a descriptive error should be returned
|
||||
require.ErrorIs(t, err, unit.ErrUnitAlreadyRegistered)
|
||||
|
||||
// Then: the unit status should not be overwritten
|
||||
u := manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusStarted, u.Status())
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("RegisterMultipleUnits", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: multiple units are registered
|
||||
unitIDs := []unit.ID{unitA, unitB, unitC}
|
||||
for _, unit := range unitIDs {
|
||||
err := manager.Register(unit)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Then: all units should be ready initially
|
||||
for _, unitID := range unitIDs {
|
||||
u := manager.Unit(unitID)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitID))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestManager_AddDependency(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("AddDependencyBetweenRegisteredUnits", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: units A and B are registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given: Unit A depends on Unit B being unit.StatusStarted
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should not be ready (depends on B)
|
||||
u := manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// Then: Unit B should still be ready (no dependencies)
|
||||
u = manager.Unit(unitB)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitB))
|
||||
|
||||
// When: Unit B is started
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should be ready, because its dependency is now in the desired state.
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
|
||||
// When: Unit B is stopped
|
||||
err = manager.UpdateStatus(unitB, unit.StatusPending)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should no longer be ready, because its dependency is not in the desired state.
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("AddDependencyByAnUnregisteredDependentUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given Unit B is registered
|
||||
err := manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given Unit A depends on Unit B being started
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
|
||||
// Then: a descriptive error communicates that the dependency cannot be added
|
||||
// because the dependent unit must be registered first.
|
||||
require.ErrorIs(t, err, unit.ErrUnitNotFound)
|
||||
})
|
||||
|
||||
t.Run("AddDependencyOnAnUnregisteredUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given unit A is registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given Unit B is not yet registered
|
||||
// And Unit A depends on Unit B being started
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
u := manager.Unit(unitB)
|
||||
assert.Equal(t, unit.StatusNotRegistered, u.Status())
|
||||
|
||||
// Then: Unit A should not be ready, because it depends on Unit B
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// When: Unit B is registered
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should still not be ready.
|
||||
// Unit B is not registered, but it has not been started as required by the dependency.
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// When: Unit B is started
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should be ready, because its dependency is now in the desired state.
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("AddDependencyCreatesACyclicDependency", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Register units
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitC)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitD)
|
||||
require.NoError(t, err)
|
||||
|
||||
// A depends on B
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
// B depends on C
|
||||
err = manager.AddDependency(unitB, unitC, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// C depends on D
|
||||
err = manager.AddDependency(unitC, unitD, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Try to make D depend on A (creates indirect cycle)
|
||||
err = manager.AddDependency(unitD, unitA, unit.StatusStarted)
|
||||
require.ErrorIs(t, err, unit.ErrCycleDetected)
|
||||
})
|
||||
}
|
||||
|
||||
func TestManager_UpdateStatus(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("UpdateStatusTriggersReadinessRecalculation", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given units A and B are registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given Unit A depends on Unit B being unit.StatusStarted
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should not be ready (depends on B)
|
||||
u := manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// When: Unit B is started
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should be ready, because its dependency is now in the desired state.
|
||||
u = manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("UpdateStatusWithUnregisteredUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given Unit A is not registered
|
||||
// When: Unit A is updated to unit.StatusStarted
|
||||
err := manager.UpdateStatus(unitA, unit.StatusStarted)
|
||||
|
||||
// Then: a descriptive error communicates that the unit must be registered first.
|
||||
require.ErrorIs(t, err, unit.ErrUnitNotFound)
|
||||
})
|
||||
|
||||
t.Run("LinearChainDependencies", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given units A, B, and C are registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitC)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create chain: A depends on B being "started", B depends on C being "completed"
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
err = manager.AddDependency(unitB, unitC, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: only Unit C should be ready (no dependencies)
|
||||
u := manager.Unit(unitC)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitC))
|
||||
|
||||
u = manager.Unit(unitB)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.False(t, manager.IsReady(unitB))
|
||||
|
||||
u = manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// When: Unit C is completed
|
||||
err = manager.UpdateStatus(unitC, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit B should be ready, because its dependency is now in the desired state.
|
||||
u = manager.Unit(unitB)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitB))
|
||||
|
||||
u = manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
u = manager.Unit(unitB)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitB))
|
||||
|
||||
// When: Unit B is started
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should be ready, because its dependency is now in the desired state.
|
||||
u = manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusPending, u.Status())
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
}
|
||||
|
||||
func TestManager_GetUnmetDependencies(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("GetUnmetDependenciesForUnitWithNoDependencies", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: Unit A is registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given: Unit A has no dependencies
|
||||
// Then: Unit A should have no unmet dependencies
|
||||
unmet, err := manager.GetUnmetDependencies(unitA)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, unmet)
|
||||
})
|
||||
|
||||
t.Run("GetUnmetDependenciesForUnitWithUnsatisfiedDependencies", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given: Unit A depends on Unit B being unit.StatusStarted
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
unmet, err := manager.GetUnmetDependencies(unitA)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, unmet, 1)
|
||||
|
||||
assert.Equal(t, unitA, unmet[0].Unit)
|
||||
assert.Equal(t, unitB, unmet[0].DependsOn)
|
||||
assert.Equal(t, unit.StatusStarted, unmet[0].RequiredStatus)
|
||||
assert.False(t, unmet[0].IsSatisfied)
|
||||
})
|
||||
|
||||
t.Run("GetUnmetDependenciesForUnitWithSatisfiedDependencies", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: Unit A and Unit B are registered
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given: Unit A depends on Unit B being unit.StatusStarted
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: Unit B is started
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should have no unmet dependencies
|
||||
unmet, err := manager.GetUnmetDependencies(unitA)
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, unmet)
|
||||
})
|
||||
|
||||
t.Run("GetUnmetDependenciesForUnregisteredUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// When: Unit A is requested
|
||||
unmet, err := manager.GetUnmetDependencies(unitA)
|
||||
|
||||
// Then: a descriptive error communicates that the unit must be registered first.
|
||||
require.ErrorIs(t, err, unit.ErrUnitNotFound)
|
||||
assert.Nil(t, unmet)
|
||||
})
|
||||
}
|
||||
|
||||
func TestManager_MultipleDependencies(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("UnitWithMultipleDependencies", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Register all units
|
||||
units := []unit.ID{unitA, unitB, unitC, unitD}
|
||||
for _, unit := range units {
|
||||
err := manager.Register(unit)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// A depends on B being unit.StatusStarted AND C being "started"
|
||||
err := manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
err = manager.AddDependency(unitA, unitC, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// A should not be ready (depends on both B and C)
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// Update B to unit.StatusStarted - A should still not be ready (needs C too)
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// Update C to "started" - A should now be ready
|
||||
err = manager.UpdateStatus(unitC, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("ComplexDependencyChain", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Register all units
|
||||
units := []unit.ID{unitA, unitB, unitC, unitD}
|
||||
for _, unit := range units {
|
||||
err := manager.Register(unit)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Create complex dependency graph:
|
||||
// A depends on B being unit.StatusStarted AND C being "started"
|
||||
err := manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
err = manager.AddDependency(unitA, unitC, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
// B depends on D being "completed"
|
||||
err = manager.AddDependency(unitB, unitD, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
// C depends on D being "completed"
|
||||
err = manager.AddDependency(unitC, unitD, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Initially only D is ready
|
||||
assert.True(t, manager.IsReady(unitD))
|
||||
assert.False(t, manager.IsReady(unitB))
|
||||
assert.False(t, manager.IsReady(unitC))
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// Update D to "completed" - B and C should become ready
|
||||
err = manager.UpdateStatus(unitD, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.True(t, manager.IsReady(unitB))
|
||||
assert.True(t, manager.IsReady(unitC))
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// Update B to unit.StatusStarted - A should still not be ready (needs C)
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// Update C to "started" - A should now be ready
|
||||
err = manager.UpdateStatus(unitC, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
|
||||
t.Run("DifferentStatusTypes", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Register units
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitC)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Given: Unit A depends on Unit B being unit.StatusStarted
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
// Given: Unit A depends on Unit C being "completed"
|
||||
err = manager.AddDependency(unitA, unitC, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
// When: Unit B is started
|
||||
err = manager.UpdateStatus(unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should not be ready, because only one of its dependencies is in the desired state.
|
||||
// It still requires Unit C to be completed.
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
|
||||
// When: Unit C is completed
|
||||
err = manager.UpdateStatus(unitC, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Then: Unit A should be ready, because both of its dependencies are in the desired state.
|
||||
assert.True(t, manager.IsReady(unitA))
|
||||
})
|
||||
}
|
||||
|
||||
func TestManager_IsReady(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("IsReadyWithUnregisteredUnit", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Given: a unit is not registered
|
||||
u := manager.Unit(unitA)
|
||||
assert.Equal(t, unit.StatusNotRegistered, u.Status())
|
||||
// Then: the unit is not ready
|
||||
assert.False(t, manager.IsReady(unitA))
|
||||
})
|
||||
}
|
||||
|
||||
func TestManager_ToDOT(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("ExportSimpleGraph", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Register units
|
||||
err := manager.Register(unitA)
|
||||
require.NoError(t, err)
|
||||
err = manager.Register(unitB)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Add dependency
|
||||
err = manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
|
||||
dot, err := manager.ExportDOT("test")
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, dot)
|
||||
assert.Contains(t, dot, "digraph")
|
||||
})
|
||||
|
||||
t.Run("ExportComplexGraph", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := unit.NewManager()
|
||||
|
||||
// Register all units
|
||||
units := []unit.ID{unitA, unitB, unitC, unitD}
|
||||
for _, unit := range units {
|
||||
err := manager.Register(unit)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// Create complex dependency graph
|
||||
// A depends on B and C, B depends on D, C depends on D
|
||||
err := manager.AddDependency(unitA, unitB, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
err = manager.AddDependency(unitA, unitC, unit.StatusStarted)
|
||||
require.NoError(t, err)
|
||||
err = manager.AddDependency(unitB, unitD, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
err = manager.AddDependency(unitC, unitD, unit.StatusComplete)
|
||||
require.NoError(t, err)
|
||||
|
||||
dot, err := manager.ExportDOT("complex")
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, dot)
|
||||
assert.Contains(t, dot, "digraph")
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user