state: unit: AssignedMachine returns a machine

worker/firewaller/firewaller_test.go

 package firewaller_test
 
 import (
         . "launchpad.net/gocheck"
         "launchpad.net/juju-core/environs"
         "launchpad.net/juju-core/environs/dummy"
         "launchpad.net/juju-core/juju/testing"
         "launchpad.net/juju-core/state"
         coretesting "launchpad.net/juju-core/testing"
         "launchpad.net/juju-core/worker/firewaller"
@@ skipping 46 matching lines @@
 
 func (s *FirewallerSuite) TestStartStop(c *C) {
         fw := firewaller.NewFirewaller(s.State)
         c.Assert(fw.Stop(), IsNil)
 }
 
 func (s *FirewallerSuite) addUnit(c *C, svc *state.Service) (*state.Unit, *state.Machine) {
         units, err := s.Conn.AddUnits(svc, 1)
         c.Assert(err, IsNil)
         u := units[0]
-        id, err := u.AssignedMachineId()
-        c.Assert(err, IsNil)
-        m, err := s.State.Machine(id)
+        m, err := u.AssignedMachine()
         c.Assert(err, IsNil)
         return u, m
 }
 
 // startInstance starts a new instance for the given machine.
 func (s *FirewallerSuite) startInstance(c *C, m *state.Machine) environs.Instance {
         inst, err := s.Conn.Environ.StartInstance(m.Id(), s.StateInfo(c), nil)
         c.Assert(err, IsNil)
         err = m.SetInstanceId(inst.Id())
         c.Assert(err, IsNil)
@@ skipping 190 matching lines @@
 }
 
 func (s *FirewallerSuite) TestFirewallerStopOnStateClose(c *C) {
         st, err := state.Open(s.StateInfo(c))
         c.Assert(err, IsNil)
         fw := firewaller.NewFirewaller(st)
         st.Close()
         c.Check(fw.Wait(), ErrorMatches, ".* zookeeper is closing")
         c.Assert(fw.Stop(), ErrorMatches, ".* zookeeper is closing")
 }