From 948bd57d6a6c8aad895be313df9fe96420a36212 Mon Sep 17 00:00:00 2001 From: Ryan Uber Date: Thu, 10 Sep 2015 12:08:08 -0700 Subject: [PATCH] agent: testing node/service maintenance using tokens --- command/agent/agent_endpoint_test.go | 18 ++++++++++++++---- command/agent/agent_test.go | 18 ++++++++++++++---- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/command/agent/agent_endpoint_test.go b/command/agent/agent_endpoint_test.go index b0f7a481c..c89977049 100644 --- a/command/agent/agent_endpoint_test.go +++ b/command/agent/agent_endpoint_test.go @@ -655,7 +655,7 @@ func TestHTTPAgent_EnableServiceMaintenance(t *testing.T) { } // Force the service into maintenance mode - req, _ := http.NewRequest("PUT", "/v1/agent/service/maintenance/test?enable=true&reason=broken", nil) + req, _ := http.NewRequest("PUT", "/v1/agent/service/maintenance/test?enable=true&reason=broken&token=mytoken", nil) resp := httptest.NewRecorder() if _, err := srv.AgentServiceMaintenance(resp, req); err != nil { t.Fatalf("err: %s", err) @@ -671,6 +671,11 @@ func TestHTTPAgent_EnableServiceMaintenance(t *testing.T) { t.Fatalf("should have registered maintenance check") } + // Ensure the token was added + if token := srv.agent.state.CheckToken(checkID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -693,7 +698,7 @@ func TestHTTPAgent_DisableServiceMaintenance(t *testing.T) { } // Force the service into maintenance mode - if err := srv.agent.EnableServiceMaintenance("test", ""); err != nil { + if err := srv.agent.EnableServiceMaintenance("test", "", ""); err != nil { t.Fatalf("err: %s", err) } @@ -749,7 +754,7 @@ func TestHTTPAgent_EnableNodeMaintenance(t *testing.T) { // Force the node into maintenance mode req, _ := http.NewRequest( - "PUT", "/v1/agent/self/maintenance?enable=true&reason=broken", nil) + "PUT", "/v1/agent/self/maintenance?enable=true&reason=broken&token=mytoken", nil) resp := httptest.NewRecorder() if _, err := srv.AgentNodeMaintenance(resp, req); err != nil { t.Fatalf("err: %s", err) @@ -764,6 +769,11 @@ func TestHTTPAgent_EnableNodeMaintenance(t *testing.T) { t.Fatalf("should have registered maintenance check") } + // Check that the token was used + if token := srv.agent.state.CheckToken(nodeMaintCheckID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -777,7 +787,7 @@ func TestHTTPAgent_DisableNodeMaintenance(t *testing.T) { defer srv.agent.Shutdown() // Force the node into maintenance mode - srv.agent.EnableNodeMaintenance("") + srv.agent.EnableNodeMaintenance("", "") // Leave maintenance mode req, _ := http.NewRequest("PUT", "/v1/agent/self/maintenance?enable=false", nil) diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go index e5fcfd09d..d735f60f3 100644 --- a/command/agent/agent_test.go +++ b/command/agent/agent_test.go @@ -1207,7 +1207,7 @@ func TestAgent_ServiceMaintenanceMode(t *testing.T) { } // Enter maintenance mode for the service - if err := agent.EnableServiceMaintenance("redis", "broken"); err != nil { + if err := agent.EnableServiceMaintenance("redis", "broken", "mytoken"); err != nil { t.Fatalf("err: %s", err) } @@ -1218,6 +1218,11 @@ func TestAgent_ServiceMaintenanceMode(t *testing.T) { t.Fatalf("should have registered critical maintenance check") } + // Check that the token was used to register the check + if token := agent.state.CheckToken(checkID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got: '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -1234,7 +1239,7 @@ func TestAgent_ServiceMaintenanceMode(t *testing.T) { } // Enter service maintenance mode without providing a reason - if err := agent.EnableServiceMaintenance("redis", ""); err != nil { + if err := agent.EnableServiceMaintenance("redis", "", ""); err != nil { t.Fatalf("err: %s", err) } @@ -1299,7 +1304,7 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) { defer agent.Shutdown() // Enter maintenance mode for the node - agent.EnableNodeMaintenance("broken") + agent.EnableNodeMaintenance("broken", "mytoken") // Make sure the critical health check was added check, ok := agent.state.Checks()[nodeMaintCheckID] @@ -1307,6 +1312,11 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) { t.Fatalf("should have registered critical node check") } + // Check that the token was used to register the check + if token := agent.state.CheckToken(nodeMaintCheckID); token != "mytoken" { + t.Fatalf("expected 'mytoken', got: '%s'", token) + } + // Ensure the reason was set in notes if check.Notes != "broken" { t.Fatalf("bad: %#v", check) @@ -1321,7 +1331,7 @@ func TestAgent_NodeMaintenanceMode(t *testing.T) { } // Enter maintenance mode without passing a reason - agent.EnableNodeMaintenance("") + agent.EnableNodeMaintenance("", "") // Make sure the check was registered with the default note check, ok = agent.state.Checks()[nodeMaintCheckID]