Continued development of the machine state.

state/state_test.go

 // launchpad.net/juju/go/state
 //
 // Copyright (c) 2011-2012 Canonical Ltd.
 package state_test
 
 import (
         "fmt"
         . "launchpad.net/gocheck"
         "launchpad.net/gozk/zookeeper"
         "launchpad.net/juju/go/charm"
@@ skipping 124 matching lines @@
 func (s StateSuite) TestGetNonExistentCharm(c *C) {
         // Check that getting a non-existent charm fails nicely.
         addDummyCharm(c, s.st)
 
         curl := charm.MustParseURL("local:anotherseries/dummy-1")
         _, err := s.st.Charm(curl)
         c.Assert(err, ErrorMatches, `charm not found: "local:anotherseries/dummy-1"`)
 }
 
 func (s StateSuite) TestAddMachine(c *C) {
-        // Check that adding machines works correctly.
         machine0, err := s.st.AddMachine()
         c.Assert(err, IsNil)
         c.Assert(machine0.Id(), Equals, 0)
         machine1, err := s.st.AddMachine()
         c.Assert(err, IsNil)
         c.Assert(machine1.Id(), Equals, 1)
 
         children, _, err := s.zkConn.Children("/machines")
         c.Assert(err, IsNil)
         sort.Strings(children)
         c.Assert(children, DeepEquals, []string{"machine-0000000000", "machine-0000000001"})
 }
 
 func (s StateSuite) TestRemoveMachine(c *C) {
-        // Check that removing a machine doesn't fail.
         machine, err := s.st.AddMachine()
         c.Assert(err, IsNil)
         _, err = s.st.AddMachine()
         c.Assert(err, IsNil)
         err = s.st.RemoveMachine(machine.Id())
         c.Assert(err, IsNil)
 
         children, _, err := s.zkConn.Children("/machines")
         c.Assert(err, IsNil)
         sort.Strings(children)
         c.Assert(children, DeepEquals, []string{"machine-0000000001"})
 
         // Removing a non-existing machine has to fail.
         err = s.st.RemoveMachine(machine.Id())
-        c.Assert(err, ErrorMatches, "machine 0 does not exist")
+        c.Assert(err, ErrorMatches, "can't remove machine 0: machine not found")
 }
 
 func (s StateSuite) TestReadMachine(c *C) {
-        // Check that reading a machine doesn't fail.
         machine, err := s.st.AddMachine()
         c.Assert(err, IsNil)
         expectedId := machine.Id()
         machine, err = s.st.Machine(expectedId)
         c.Assert(err, IsNil)
         c.Assert(machine.Id(), Equals, expectedId)
 }
 
 func (s StateSuite) TestReadNonExistentMachine(c *C) {
-        // Check that reading a non-existent machine fails nicely.
         _, err := s.st.Machine(0)
         c.Assert(err, ErrorMatches, "machine 0 not found")
 
         _, err = s.st.AddMachine()
         c.Assert(err, IsNil)
         _, err = s.st.Machine(1)
         c.Assert(err, ErrorMatches, "machine 1 not found")
 }
 
 func (s StateSuite) TestAllMachines(c *C) {
-        // Check that reading all machines works correctly.
         machines, err := s.st.AllMachines()
         c.Assert(err, IsNil)
         c.Assert(len(machines), Equals, 0)
 
         _, err = s.st.AddMachine()
         c.Assert(err, IsNil)
         machines, err = s.st.AllMachines()
         c.Assert(err, IsNil)
         c.Assert(len(machines), Equals, 1)
 
         _, err = s.st.AddMachine()
         c.Assert(err, IsNil)
         machines, err = s.st.AllMachines()
         c.Assert(err, IsNil)
         c.Assert(len(machines), Equals, 2)
 }
 
 func (s StateSuite) TestAddService(c *C) {
-        // Check that adding services works correctly.
         dummy, curl := addDummyCharm(c, s.st)
         wordpress, err := s.st.AddService("wordpress", dummy)
         c.Assert(err, IsNil)
         c.Assert(wordpress.Name(), Equals, "wordpress")
         mysql, err := s.st.AddService("mysql", dummy)
         c.Assert(err, IsNil)
         c.Assert(mysql.Name(), Equals, "mysql")
 
         // Check that retrieving the new created services works correctly.
         wordpress, err = s.st.Service("wordpress")
         c.Assert(err, IsNil)
         c.Assert(wordpress.Name(), Equals, "wordpress")
         url, err := wordpress.CharmURL()
         c.Assert(err, IsNil)
         c.Assert(url.String(), Equals, curl.String())
         mysql, err = s.st.Service("mysql")
         c.Assert(err, IsNil)
         c.Assert(mysql.Name(), Equals, "mysql")
         url, err = mysql.CharmURL()
         c.Assert(err, IsNil)
         c.Assert(url.String(), Equals, curl.String())
 }
 
 func (s StateSuite) TestRemoveService(c *C) {
         dummy, _ := addDummyCharm(c, s.st)
         service, err := s.st.AddService("wordpress", dummy)
         c.Assert(err, IsNil)
 
-        // Check that removing the service works correctly.
         err = s.st.RemoveService(service)
         c.Assert(err, IsNil)
         service, err = s.st.Service("wordpress")
         c.Assert(err, ErrorMatches, `service with name "wordpress" not found`)
 }
 
 func (s StateSuite) TestReadNonExistentService(c *C) {
-        // Check that reading a non-existent service fails nicely.
         _, err := s.st.Service("pressword")
         c.Assert(err, ErrorMatches, `service with name "pressword" not found`)
 }
 
 func (s StateSuite) TestAllServices(c *C) {
-        // Check without existing services.
         services, err := s.st.AllServices()
         c.Assert(err, IsNil)
         c.Assert(len(services), Equals, 0)
 
         // Check that after adding services the result is ok.
         dummy, _ := addDummyCharm(c, s.st)
         _, err = s.st.AddService("wordpress", dummy)
         c.Assert(err, IsNil)
         services, err = s.st.AllServices()
         c.Assert(err, IsNil)
@@ skipping 416 matching lines @@
         err = mysqlUnit.AssignToMachine(mysqlMachine)
         c.Assert(err, IsNil)
 
         wordpressService, err := s.st.AddService("wordpress", dummy)
         c.Assert(err, IsNil)
         wordpressUnit, err := wordpressService.AddUnit()
         c.Assert(err, IsNil)
 
         _, err = wordpressUnit.AssignToUnusedMachine()
         c.Assert(err, ErrorMatches, "no unused machine found")
+}
+
+func (s StateSuite) TestGetSetClearUnitUpgrate(c *C) {
+        // Check that setting and clearing an upgrade flag on a unit works.
+        dummy, _ := addDummyCharm(c, s.st)
+        wordpress, err := s.st.AddService("wordpress", dummy)
+        c.Assert(err, IsNil)
+        unit, err := wordpress.AddUnit()
+        c.Assert(err, IsNil)
+
+        // Defaults to false.
+        upgrade, err := unit.NeedsUpgrade()
+        c.Assert(err, IsNil)
+        c.Assert(upgrade, Equals, false)
+
+        // Can be set.
+        err = unit.SetNeedsUpgrade()
+        c.Assert(err, IsNil)
+        upgrade, err = unit.NeedsUpgrade()
+        c.Assert(err, IsNil)
+        c.Assert(upgrade, Equals, true)
+
+        // Can be set multiple times.
+        err = unit.SetNeedsUpgrade()
+        c.Assert(err, IsNil)
+        upgrade, err = unit.NeedsUpgrade()
+        c.Assert(err, IsNil)
+        c.Assert(upgrade, Equals, true)
+
+        // Can be cleared.
+        err = unit.ClearNeedsUpgrade()
+        c.Assert(err, IsNil)
+        upgrade, err = unit.NeedsUpgrade()
+        c.Assert(err, IsNil)
+        c.Assert(upgrade, Equals, false)
+
+        // Can be cleared multiple times
+        err = unit.ClearNeedsUpgrade()
+        c.Assert(err, IsNil)
+        upgrade, err = unit.NeedsUpgrade()
+        c.Assert(err, IsNil)
+        c.Assert(upgrade, Equals, false)
+}
+
+func (s StateSuite) TestGetSetClearResolved(c *C) {
+        // Check that setting and clearing the resolved setting on a unit works.
+        dummy, _ := addDummyCharm(c, s.st)
+        wordpress, err := s.st.AddService("wordpress", dummy)
+        c.Assert(err, IsNil)
+        unit, err := wordpress.AddUnit()
+        c.Assert(err, IsNil)
+
+        setting, err := unit.Resolved()
+        c.Assert(err, IsNil)
+        c.Assert(setting, Equals, state.ResolvedNone)
+
+        err = unit.SetResolved(state.ResolvedNoHooks)
+        c.Assert(err, IsNil)
+        err = unit.SetResolved(state.ResolvedNoHooks)
+        c.Assert(err, ErrorMatches, `unit "wordpress/0" resolved flag already set`)
+        retry, err := unit.Resolved()
+        c.Assert(err, IsNil)
+        c.Assert(retry, Equals, state.ResolvedNoHooks)
+
+        err = unit.ClearResolved()
+        c.Assert(err, IsNil)
+        setting, err = unit.Resolved()
+        c.Assert(err, IsNil)
+        c.Assert(setting, Equals, state.ResolvedNone)
+        err = unit.ClearResolved()
+        c.Assert(err, IsNil)
+
+        err = unit.SetResolved(state.ResolvedMode(999))
+        c.Assert(err, ErrorMatches, `invalid error resolution mode: 999`)
+}
+
+func (s StateSuite) TestGetOpenPorts(c *C) {
+        // Check that changes to the open ports of units work porperly.
+        dummy, _ := addDummyCharm(c, s.st)
+        wordpress, err := s.st.AddService("wordpress", dummy)
+        c.Assert(err, IsNil)
+        unit, err := wordpress.AddUnit()
+        c.Assert(err, IsNil)
+
+        // Verify no open ports before activity.
+        open, err := unit.OpenPorts()
+        c.Assert(err, IsNil)
+        c.Assert(open, HasLen, 0)
+
+        // Now open and close port.
+        err = unit.OpenPort("tcp", 80)
+        c.Assert(err, IsNil)
+        open, err = unit.OpenPorts()
+        c.Assert(err, IsNil)
+        c.Assert(open, DeepEquals, []state.Port{
+                {"tcp", 80},
+        })
+
+        err = unit.OpenPort("udp", 53)
+        c.Assert(err, IsNil)
+        open, err = unit.OpenPorts()
+        c.Assert(err, IsNil)
+        c.Assert(open, DeepEquals, []state.Port{
+                {"tcp", 80},
+                {"udp", 53},
+        })
+
+        err = unit.OpenPort("tcp", 53)
+        c.Assert(err, IsNil)
+        open, err = unit.OpenPorts()
+        c.Assert(err, IsNil)
+        c.Assert(open, DeepEquals, []state.Port{
+                {"tcp", 80},
+                {"udp", 53},
+                {"tcp", 53},
+        })
+
+        err = unit.OpenPort("tcp", 443)
+        c.Assert(err, IsNil)
+        open, err = unit.OpenPorts()
+        c.Assert(err, IsNil)
+        c.Assert(open, DeepEquals, []state.Port{
+                {"tcp", 80},
+                {"udp", 53},
+                {"tcp", 53},
+                {"tcp", 443},
+        })
+
+        err = unit.ClosePort("tcp", 80)
+        c.Assert(err, IsNil)
+        open, err = unit.OpenPorts()
+        c.Assert(err, IsNil)
+        c.Assert(open, DeepEquals, []state.Port{
+                {"udp", 53},
+                {"tcp", 53},
+                {"tcp", 443},
+        })
 }
 
 // zkRemoveTree recursively removes a tree.
 func zkRemoveTree(zk *zookeeper.Conn, path string) error {
         // First recursively delete the children.
         children, _, err := zk.Children(path)
         if err != nil {
                 return err
         }
         for _, child := range children {
                 if err = zkRemoveTree(zk, fmt.Sprintf("%s/%s", path, child)); err != nil {
                         return err
                 }
         }
         // Now delete the path itself.
         return zk.Delete(path, -1)
 }