Repository: hashicorp/nomad
File: client/allocrunner/group_service_hook_test.go (Go)

Code before:
			logger,
			consulMockClient,
			regMock.NewServiceRegistrationHandler(logger))

		shutDownCtx, cancel := context.WithTimeout(context.Background(), delay*2)
		defer cancel()

		h := newGroupServiceHook(groupServiceHookConfig{
			alloc:             alloc,
			serviceRegWrapper: regWrapper,
			shutdownDelayCtx:  shutDownCtx,
			restarter:         agentconsul.NoopRestarter(),
			taskEnvBuilder:    taskenv.NewBuilder(mock.Node(), alloc, nil, alloc.Job.Region),
			logger:            logger,
			hookResources:     cstructs.NewAllocHookResources(),
		})

		successChan := make(chan struct{}, 1)
		go func() {
			before := time.Now()
			h.PreKill()
			after := time.Now()
			// we respected the shutdown_delay and did not return immediately
			if after.Sub(before) < delay {
				t.Fail()
			}
			successChan <- struct{}{}
		}()

		select {
		case <-shutDownCtx.Done():
			// we didn't wait longer than shutdown delay and get killed by the shutdownCtx
			t.Fail()
		case <-successChan:
		}
	})

	t.Run("returns immediately when no shutdown delay", func(t *testing.T) {
		alloc := mock.Alloc()
		alloc.Job.TaskGroups[0].Networks = []*structs.NetworkResource{}
		tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
		tg.Services = []*structs.Service{
			{
				Name:      "testconnect",
				PortLabel: "9999",
				Connect: &structs.ConsulConnect{
					SidecarService: &structs.ConsulSidecarService{},
				},
			},
		}


Review comment (bug):
The stdlib has a nice little helper function for this. And I think the comment is backwards?

```suggestion
			// we did not respect the shutdown_delay and returned immediately
			if time.Since(before) < delay {
				t.Fail()
			}
```

Code after:
			logger,
			consulMockClient,
			regMock.NewServiceRegistrationHandler(logger))

		shutDownCtx, cancel := context.WithTimeout(context.Background(), delay*2)
		defer cancel()

		h := newGroupServiceHook(groupServiceHookConfig{
			alloc:             alloc,
			serviceRegWrapper: regWrapper,
			shutdownDelayCtx:  shutDownCtx,
			restarter:         agentconsul.NoopRestarter(),
			taskEnvBuilder:    taskenv.NewBuilder(mock.Node(), alloc, nil, alloc.Job.Region),
			logger:            logger,
			hookResources:     cstructs.NewAllocHookResources(),
		})

		successChan := make(chan struct{}, 1)
		go func() {
			before := time.Now()
			h.PreKill()
			// we did not respect the shutdown_delay and returned immediately
			if time.Since(before) < delay {
				t.Fail()
			}
			successChan <- struct{}{}
		}()

		select {
		case <-shutDownCtx.Done():
			// we didn't wait longer than shutdown delay and get killed by the shutdownCtx
			t.Fail()
		case <-successChan:
		}
	})

	t.Run("returns immediately when no shutdown delay", func(t *testing.T) {
		alloc := mock.Alloc()
		alloc.Job.TaskGroups[0].Networks = []*structs.NetworkResource{}
		tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
		tg.Services = []*structs.Service{
			{
				Name:      "testconnect",
				PortLabel: "9999",
				Connect: &structs.ConsulConnect{
					SidecarService: &structs.ConsulSidecarService{},
				},
			},
		}

		consulMockClient := regMock.NewServiceRegistrationHandler(logger)