Skip to content

Commit 1dd8983

Browse files
tgrosstgross
committed
lifecycle: unit test for lifecycle task behavior on restarts
1 parent 6dcada4 commit 1dd8983

File tree

1 file changed

+253
-0
lines changed

1 file changed

+253
-0
lines changed

‎client/allocrunner/alloc_runner_test.go

+253
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"time"
1010

1111
"github.com/hashicorp/consul/api"
12+
multierror "github.com/hashicorp/go-multierror"
1213
"github.com/hashicorp/nomad/client/allochealth"
1314
"github.com/hashicorp/nomad/client/allocwatcher"
1415
cconsul "github.com/hashicorp/nomad/client/consul"
@@ -1568,3 +1569,255 @@ func TestAllocRunner_PersistState_Destroyed(t *testing.T) {
15681569
require.NoError(t, err)
15691570
require.Nil(t, ts)
15701571
}
1572+
1573+
func TestAllocRunner_Restart(t *testing.T) {
1574+
1575+
alloc := mock.LifecycleAlloc()
1576+
alloc.Job.Type = structs.JobTypeService
1577+
1578+
// helper for constructing lifecycle tasks. The LifecycleAlloc and
1579+
// LifecycleJob don't have *all* the possible lifecycle options, and the
1580+
// allocrunner needs the resource fields filled out.
1581+
//
1582+
// TODO: should this larger set just get moved to the mock? We'll
1583+
// want to check who else is using it
1584+
addTask := func(name, runFor string, lc *structs.TaskLifecycleConfig) {
1585+
alloc.Job.TaskGroups[0].Tasks = append(alloc.Job.TaskGroups[0].Tasks,
1586+
&structs.Task{
1587+
Name: name,
1588+
Driver: "mock_driver",
1589+
Config: map[string]interface{}{"run_for": runFor},
1590+
Lifecycle: lc,
1591+
LogConfig: structs.DefaultLogConfig(),
1592+
Resources: &structs.Resources{CPU: 100, MemoryMB: 256},
1593+
},
1594+
)
1595+
alloc.TaskResources[name] = &structs.Resources{CPU: 100, MemoryMB: 256}
1596+
alloc.AllocatedResources.Tasks[name] = &structs.AllocatedTaskResources{
1597+
Cpu: structs.AllocatedCpuResources{CpuShares: 100},
1598+
Memory: structs.AllocatedMemoryResources{MemoryMB: 256},
1599+
}
1600+
}
1601+
1602+
addTask("main", "100s", nil)
1603+
addTask("prestart-oneshot", "1s", &structs.TaskLifecycleConfig{
1604+
Hook: structs.TaskLifecycleHookPrestart,
1605+
Sidecar: false,
1606+
})
1607+
addTask("prestart-sidecar", "100s", &structs.TaskLifecycleConfig{
1608+
Hook: structs.TaskLifecycleHookPrestart,
1609+
Sidecar: true,
1610+
})
1611+
addTask("poststart-oneshot", "1s", &structs.TaskLifecycleConfig{
1612+
Hook: structs.TaskLifecycleHookPrestart,
1613+
Sidecar: false,
1614+
})
1615+
addTask("poststart-sidecar", "100s", &structs.TaskLifecycleConfig{
1616+
Hook: structs.TaskLifecycleHookPrestart,
1617+
Sidecar: true,
1618+
})
1619+
addTask("poststop", "1s", &structs.TaskLifecycleConfig{
1620+
Hook: structs.TaskLifecycleHookPoststop,
1621+
})
1622+
1623+
testCases := []struct {
1624+
name string
1625+
action func(*allocRunner, *structs.Allocation) error
1626+
expectedAfter map[string]structs.TaskState
1627+
}{
1628+
{
1629+
name: "restart entire allocation",
1630+
action: func(ar *allocRunner, alloc *structs.Allocation) error {
1631+
// TODO: is there useful test info we could get from a TaskEvent here?
1632+
ar.RestartAll(&structs.TaskEvent{})
1633+
return nil
1634+
},
1635+
expectedAfter: map[string]structs.TaskState{
1636+
"main": structs.TaskState{
1637+
State: structs.TaskStateRunning,
1638+
Restarts: 1,
1639+
},
1640+
"prestart-oneshot": structs.TaskState{
1641+
State: structs.TaskStateDead,
1642+
Restarts: 0, // TODO: is this true?
1643+
},
1644+
"prestart-sidecar": structs.TaskState{
1645+
State: structs.TaskStateRunning,
1646+
Restarts: 1, // TODO: is this true?
1647+
},
1648+
"poststart-oneshot": structs.TaskState{
1649+
State: structs.TaskStateDead,
1650+
Restarts: 0, // TODO: is this true?
1651+
},
1652+
"poststart-sidecar": structs.TaskState{
1653+
State: structs.TaskStateRunning,
1654+
Restarts: 1, // TODO: is this true?
1655+
},
1656+
"poststop": structs.TaskState{
1657+
State: structs.TaskStatePending,
1658+
Restarts: 0,
1659+
},
1660+
},
1661+
},
1662+
1663+
{
1664+
name: "restart main task",
1665+
action: func(ar *allocRunner, alloc *structs.Allocation) error {
1666+
ar.RestartTask("main", &structs.TaskEvent{})
1667+
return nil
1668+
},
1669+
expectedAfter: map[string]structs.TaskState{
1670+
"main": structs.TaskState{
1671+
State: structs.TaskStateRunning,
1672+
Restarts: 1,
1673+
},
1674+
"prestart-oneshot": structs.TaskState{
1675+
State: structs.TaskStateDead,
1676+
Restarts: 0, // TODO: is this true?
1677+
},
1678+
"prestart-sidecar": structs.TaskState{
1679+
State: structs.TaskStateRunning,
1680+
Restarts: 0, // TODO: is this true?
1681+
},
1682+
"poststart-oneshot": structs.TaskState{
1683+
State: structs.TaskStateDead,
1684+
Restarts: 0, // TODO: is this true?
1685+
},
1686+
"poststart-sidecar": structs.TaskState{
1687+
State: structs.TaskStateRunning,
1688+
Restarts: 1, // TODO: is this true?
1689+
},
1690+
"poststop": structs.TaskState{
1691+
State: structs.TaskStatePending,
1692+
Restarts: 0,
1693+
},
1694+
},
1695+
},
1696+
1697+
{
1698+
name: "restart prestart-sidecar task",
1699+
action: func(ar *allocRunner, alloc *structs.Allocation) error {
1700+
ar.RestartTask("prestart-sidecar", &structs.TaskEvent{})
1701+
return nil
1702+
},
1703+
expectedAfter: map[string]structs.TaskState{
1704+
"main": structs.TaskState{
1705+
State: structs.TaskStateRunning,
1706+
Restarts: 0, // TODO: is this true?
1707+
},
1708+
"prestart-oneshot": structs.TaskState{
1709+
State: structs.TaskStateDead,
1710+
Restarts: 0, // TODO: is this true?
1711+
},
1712+
"prestart-sidecar": structs.TaskState{
1713+
State: structs.TaskStateRunning,
1714+
Restarts: 1,
1715+
},
1716+
"poststart-oneshot": structs.TaskState{
1717+
State: structs.TaskStateDead,
1718+
Restarts: 0, // TODO: is this true?
1719+
},
1720+
"poststart-sidecar": structs.TaskState{
1721+
State: structs.TaskStateRunning,
1722+
Restarts: 0, // TODO: is this true?
1723+
},
1724+
"poststop": structs.TaskState{
1725+
State: structs.TaskStatePending,
1726+
Restarts: 0,
1727+
},
1728+
},
1729+
},
1730+
}
1731+
1732+
for _, tc := range testCases {
1733+
tc := tc
1734+
t.Run(tc.name, func(t *testing.T) {
1735+
t.Parallel()
1736+
require := require.New(t)
1737+
1738+
alloc := alloc.Copy()
1739+
1740+
conf, cleanup := testAllocRunnerConfig(t, alloc)
1741+
defer cleanup()
1742+
ar, err := NewAllocRunner(conf)
1743+
require.NoError(err)
1744+
defer destroy(ar)
1745+
go ar.Run()
1746+
1747+
upd := conf.StateUpdater.(*MockStateUpdater)
1748+
1749+
// assert our "before" states
1750+
testutil.WaitForResult(func() (bool, error) {
1751+
last := upd.Last()
1752+
if last == nil {
1753+
return false, fmt.Errorf("no update")
1754+
}
1755+
if last.ClientStatus != structs.AllocClientStatusRunning {
1756+
return false, fmt.Errorf(
1757+
"expected alloc to be running not %s", last.ClientStatus)
1758+
}
1759+
var errs *multierror.Error
1760+
1761+
expectedBefore := map[string]string{
1762+
"main": structs.TaskStateRunning,
1763+
"prestart-oneshot": structs.TaskStateDead,
1764+
"prestart-sidecar": structs.TaskStateRunning,
1765+
"poststart-oneshot": structs.TaskStateDead,
1766+
"poststart-sidecar": structs.TaskStateRunning,
1767+
"poststop": structs.TaskStatePending,
1768+
}
1769+
1770+
for task, expected := range expectedBefore {
1771+
got := last.TaskStates[task]
1772+
if got.State != expected.State {
1773+
errs = multierror.Append(errs, fmt.Errorf(
1774+
"expected initial state of task %q to be %q not %q",
1775+
task, expected.State, got.State))
1776+
}
1777+
if got.Restarts != 0 {
1778+
errs = multierror.Append(errs, fmt.Errorf(
1779+
"expected no initial restarts of task %q, not %q",
1780+
task, expected.State, got.State))
1781+
}
1782+
}
1783+
if errs.ErrorOrNil() != nil {
1784+
return false, errs.ErrorOrNil()
1785+
}
1786+
return true, nil
1787+
}, func(err error) {
1788+
require.NoError(err, "error waiting for initial state")
1789+
})
1790+
1791+
// perform the action
1792+
tc.action(ar, alloc.Copy())
1793+
1794+
// assert our "after" states
1795+
testutil.WaitForResult(func() (bool, error) {
1796+
last := upd.Last()
1797+
if last == nil {
1798+
return false, fmt.Errorf("no update")
1799+
}
1800+
var errs *multierror.Error
1801+
for task, expected := range tc.expectedAfter {
1802+
got := last.TaskStates[task]
1803+
if got.State != expected.State {
1804+
errs = multierror.Append(errs, fmt.Errorf(
1805+
"expected final state of task %q to be %q not %q",
1806+
task, expected.State, got.State))
1807+
}
1808+
if got.Restarts != expected.Restarts {
1809+
errs = multierror.Append(errs, fmt.Errorf(
1810+
"expected final restarts of task %q to be %q not %q",
1811+
task, expected.State, got.State))
1812+
}
1813+
}
1814+
if errs.ErrorOrNil() != nil {
1815+
return false, errs.ErrorOrNil()
1816+
}
1817+
return true, nil
1818+
}, func(err error) {
1819+
require.NoError(err, "error waiting for final state")
1820+
})
1821+
})
1822+
}
1823+
}

0 commit comments

Comments
 (0)