func TestTaskManagerRunTask(t *testing.T) { resolved := false var serverURL string var handler = func(w http.ResponseWriter, r *http.Request) { if strings.Contains(r.URL.Path, "/artifacts/public/logs/live_backing.log") { json.NewEncoder(w).Encode(&queue.S3ArtifactResponse{ PutURL: serverURL, }) return } if strings.Contains(r.URL.Path, "/artifacts/public/logs/live.log") { json.NewEncoder(w).Encode(&queue.RedirectArtifactResponse{}) return } if strings.Contains(r.URL.Path, "/task/abc/runs/1/completed") { resolved = true w.Header().Set("Content-Type", "application/json; charset=UTF-8") json.NewEncoder(w).Encode(&queue.TaskStatusResponse{}) } } s := httptest.NewServer(http.HandlerFunc(handler)) serverURL = s.URL defer s.Close() tempPath := filepath.Join(os.TempDir(), slugid.Nice()) tempStorage, err := runtime.NewTemporaryStorage(tempPath) if err != nil { t.Fatal(err) } localServer, err := webhookserver.NewLocalServer( []byte{127, 0, 0, 1}, 60000, "", 0, "example.com", "", "", "", 10*time.Minute, ) if err != nil { t.Error(err) } gc := &gc.GarbageCollector{} environment := &runtime.Environment{ GarbageCollector: gc, TemporaryStorage: tempStorage, WebHookServer: localServer, } engineProvider := engines.Engines()["mock"] engine, err := engineProvider.NewEngine(engines.EngineOptions{ Environment: environment, Log: logger.WithField("engine", "mock"), }) if err != nil { t.Fatal(err.Error()) } cfg := &configType{ QueueBaseURL: serverURL, } tm, err := newTaskManager(cfg, engine, MockPlugin{}, environment, logger.WithField("test", "TestTaskManagerRunTask"), gc) if err != nil { t.Fatal(err) } claim := &taskClaim{ taskID: "abc", runID: 1, taskClaim: &queue.TaskClaimResponse{ Credentials: struct { AccessToken string `json:"accessToken"` Certificate string `json:"certificate"` ClientID string `json:"clientId"` }{ AccessToken: "123", ClientID: "abc", Certificate: "", }, TakenUntil: tcclient.Time(time.Now().Add(time.Minute * 5)), }, definition: &queue.TaskDefinitionResponse{ Payload: []byte(`{"delay": 1,"function": "write-log","argument": "Hello World"}`), }, } tm.run(claim) assert.True(t, resolved, "Task was not resolved") }
// New will create a worker given configuration matching the schema from // ConfigSchema(). The log parameter is optional and if nil is given a default // logrus logger will be used. func New(config interface{}, log *logrus.Logger) (*Worker, error) { // Validate and map configuration to c var c configType if err := schematypes.MustMap(ConfigSchema(), config, &c); err != nil { return nil, fmt.Errorf("Invalid configuration: %s", err) } // Create temporary folder err := os.RemoveAll(c.TemporaryFolder) if err != nil { return nil, fmt.Errorf("Failed to remove temporaryFolder: %s, error: %s", c.TemporaryFolder, err) } tempStorage, err := runtime.NewTemporaryStorage(c.TemporaryFolder) if err != nil { return nil, fmt.Errorf("Failed to create temporary folder, error: %s", err) } // Create logger if log == nil { log = logrus.New() } log.Level, _ = logrus.ParseLevel(c.LogLevel) // Setup WebHookServer localServer, err := webhookserver.NewLocalServer( net.ParseIP(c.ServerIP), c.ServerPort, c.NetworkInterface, c.ExposedPort, c.DNSDomain, c.DNSSecret, c.TLSCertificate, c.TLSKey, time.Duration(c.MaxLifeCycle)*time.Second, ) if err != nil { return nil, err } // Create environment gc := gc.New(c.TemporaryFolder, c.MinimumDiskSpace, c.MinimumMemory) env := &runtime.Environment{ GarbageCollector: gc, Log: log, TemporaryStorage: tempStorage, WebHookServer: localServer, } // Ensure that engine confiuguration was provided for the engine selected if _, ok := c.Engines[c.Engine]; !ok { return nil, fmt.Errorf("Invalid configuration: The key 'engines.%s' must "+ "be specified when engine '%s' is selected", c.Engine, c.Engine) } // Find engine provider (schema should ensure it exists) provider := engines.Engines()[c.Engine] engine, err := provider.NewEngine(engines.EngineOptions{ Environment: env, Log: env.Log.WithField("engine", c.Engine), Config: c.Engines[c.Engine], }) if err != nil { return nil, fmt.Errorf("Engine initialization failed, error: %s", err) } // Initialize plugin manager pm, err := plugins.NewPluginManager(plugins.PluginOptions{ Environment: env, Engine: engine, Log: env.Log.WithField("plugin", "plugin-manager"), Config: c.Plugins, }) if err != nil { return nil, fmt.Errorf("Plugin initialization failed, error: %s", err) } tm, err := newTaskManager( &c, engine, pm, env, env.Log.WithField("component", "task-manager"), gc, ) if err != nil { return nil, err } return &Worker{ log: env.Log.WithField("component", "worker"), tm: tm, sm: runtime.NewShutdownManager("local"), env: env, server: localServer, done: make(chan struct{}), }, nil }
func TestWorkerShutdown(t *testing.T) { var resCount int32 var serverURL string var handler = func(w http.ResponseWriter, r *http.Request) { if strings.Contains(r.URL.Path, "/artifacts/public/logs/live_backing.log") { json.NewEncoder(w).Encode(&queue.S3ArtifactResponse{ PutURL: serverURL, }) return } if strings.Contains(r.URL.Path, "/artifacts/public/logs/live.log") { json.NewEncoder(w).Encode(&queue.RedirectArtifactResponse{}) return } if strings.Contains(r.URL.Path, "exception") { var exception queue.TaskExceptionRequest err := json.NewDecoder(r.Body).Decode(&exception) // Ignore errors for now if err != nil { return } assert.Equal(t, "worker-shutdown", exception.Reason) atomic.AddInt32(&resCount, 1) w.Header().Set("Content-Type", "application/json; charset=UTF-8") json.NewEncoder(w).Encode(&queue.TaskStatusResponse{}) } } s := httptest.NewServer(http.HandlerFunc(handler)) serverURL = s.URL defer s.Close() tempPath := filepath.Join(os.TempDir(), slugid.Nice()) tempStorage, err := runtime.NewTemporaryStorage(tempPath) if err != nil { t.Fatal(err) } localServer, err := webhookserver.NewLocalServer( []byte{127, 0, 0, 1}, 60000, "", 0, "example.com", "", "", "", 10*time.Minute, ) if err != nil { t.Error(err) } gc := &gc.GarbageCollector{} environment := &runtime.Environment{ GarbageCollector: gc, TemporaryStorage: tempStorage, WebHookServer: localServer, } engineProvider := engines.Engines()["mock"] engine, err := engineProvider.NewEngine(engines.EngineOptions{ Environment: environment, Log: logger.WithField("engine", "mock"), }) if err != nil { t.Fatal(err.Error()) } cfg := &configType{ QueueBaseURL: serverURL, } tm, err := newTaskManager(cfg, engine, MockPlugin{}, environment, logger.WithField("test", "TestRunTask"), gc) if err != nil { t.Fatal(err) } claims := []*taskClaim{ &taskClaim{ taskID: "abc", runID: 1, definition: &queue.TaskDefinitionResponse{ Payload: []byte(`{"delay": 5000,"function": "write-log","argument": "Hello World"}`), }, taskClaim: &queue.TaskClaimResponse{ Credentials: struct { AccessToken string `json:"accessToken"` Certificate string `json:"certificate"` ClientID string `json:"clientId"` }{ AccessToken: "123", ClientID: "abc", Certificate: "", }, TakenUntil: tcclient.Time(time.Now().Add(time.Minute * 5)), }, }, &taskClaim{ taskID: "def", runID: 0, definition: &queue.TaskDefinitionResponse{ Payload: []byte(`{"delay": 5000,"function": "write-log","argument": "Hello World"}`), }, taskClaim: &queue.TaskClaimResponse{ Credentials: struct { AccessToken string `json:"accessToken"` Certificate string `json:"certificate"` ClientID string `json:"clientId"` }{ AccessToken: "123", ClientID: "abc", Certificate: "", }, TakenUntil: tcclient.Time(time.Now().Add(time.Minute * 5)), }, }, } var wg sync.WaitGroup wg.Add(2) go func() { for _, c := range claims { go func(claim *taskClaim) { tm.run(claim) wg.Done() }(c) } }() time.Sleep(500 * time.Millisecond) assert.Equal(t, len(tm.RunningTasks()), 2) close(tm.done) tm.Stop() wg.Wait() assert.Equal(t, 0, len(tm.RunningTasks())) assert.Equal(t, int32(2), atomic.LoadInt32(&resCount)) }