cmd/jujud: fix race in test

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 42 matching lines @@
 func (s *FirewallerSuite) SetUpTest(c *C) {
         s.JujuConnSuite.SetUpTest(c)
         s.charm = s.AddTestingCharm(c, "dummy")
 }
 
 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)
+        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)
+
+        // Check that instance has actually been started
+        config, err := s.State.EnvironConfig()
+        c.Assert(err, IsNil)
+        env, err := environs.NewFromAttrs(config.Map())

[TheMue, 2012/08/30 16:00:44]

After merge you can use environs.New(config) here.

[rog, 2012/08/31 05:48:01]

removed this logic. it's not testing anything that hasn't been tested elsewhere.

+        _, err = env.Instances([]string{inst.Id()})
+        c.Assert(err, IsNil)
+        return inst
+}
+
 func (s *FirewallerSuite) TestNotExposedService(c *C) {
         fw := firewaller.NewFirewaller(s.State)
         defer func() { c.Assert(fw.Stop(), IsNil) }()
 
-        m, err := s.State.AddMachine()
+        svc, err := s.Conn.AddService("wordpress", s.charm)
         c.Assert(err, IsNil)
-        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)
+        u, m := s.addUnit(c, svc)
+        inst := s.startInstance(c, m)
 
-        svc, err := s.State.AddService("wordpress", s.charm)
-        c.Assert(err, IsNil)
-        u, err := svc.AddUnit()
-        c.Assert(err, IsNil)
-        err = u.AssignToMachine(m)
-        c.Assert(err, IsNil)
         err = u.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
         err = u.OpenPort("tcp", 8080)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), nil)
 
         err = u.ClosePort("tcp", 80)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), nil)
 }
 
 func (s *FirewallerSuite) TestExposedService(c *C) {
         fw := firewaller.NewFirewaller(s.State)
         defer func() { c.Assert(fw.Stop(), IsNil) }()
 
-        m, err := s.State.AddMachine()
-        c.Assert(err, IsNil)
-        inst, err := s.Conn.Environ.StartInstance(m.Id(), s.StateInfo(c), nil)
-        c.Assert(err, IsNil)
-        err = m.SetInstanceId(inst.Id())
+        svc, err := s.Conn.AddService("wordpress", s.charm)
         c.Assert(err, IsNil)
 
-        config, err := s.State.EnvironConfig()
-        c.Assert(err, IsNil)
-        env, err := environs.NewFromAttrs(config.Map())
-        _, err = env.Instances([]string{inst.Id()})
-        c.Assert(err, IsNil)
-        c.Logf("got instance fine, id %q", inst.Id())
-
-        svc, err := s.State.AddService("wordpress", s.charm)
-        c.Assert(err, IsNil)
         err = svc.SetExposed()
         c.Assert(err, IsNil)
-        u, err := svc.AddUnit()
-        c.Assert(err, IsNil)
-        err = u.AssignToMachine(m)
-        c.Assert(err, IsNil)
+        u, m := s.addUnit(c, svc)
+        inst := s.startInstance(c, m)
+
         err = u.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
         err = u.OpenPort("tcp", 8080)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), []state.Port{{"tcp", 80}, {"tcp", 8080}})
 
         err = u.ClosePort("tcp", 80)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), []state.Port{{"tcp", 8080}})
 }
 
+func (s *FirewallerSuite) TestMachineWithoutInstanceId(c *C) {
+        fw := firewaller.NewFirewaller(s.State)
+        defer func() { c.Assert(fw.Stop(), IsNil) }()
+
+        svc, err := s.Conn.AddService("wordpress", s.charm)
+        c.Assert(err, IsNil)
+        err = svc.SetExposed()
+        c.Assert(err, IsNil)
+        // add a unit but don't start its instance yet.
+        u1, m1 := s.addUnit(c, svc)
+
+        // add another unit and start its instance, so that
+        // we're sure the firewaller has seen the first instance.
+        u2, m2 := s.addUnit(c, svc)
+        inst2 := s.startInstance(c, m2)
+        err = u2.OpenPort("tcp", 80)
+        c.Assert(err, IsNil)
+        assertPorts(c, inst2, m2.Id(), []state.Port{{"tcp", 80}})
+
+        inst1 := s.startInstance(c, m1)
+        err = u1.OpenPort("tcp", 8080)
+        c.Assert(err, IsNil)
+        assertPorts(c, inst1, m1.Id(), []state.Port{{"tcp", 8080}})
+}
+
 func (s *FirewallerSuite) TestMultipleUnits(c *C) {
         fw := firewaller.NewFirewaller(s.State)
         defer func() { c.Assert(fw.Stop(), IsNil) }()
 
-        m1, err := s.State.AddMachine()
-        c.Assert(err, IsNil)
-
-        inst1, err := s.Conn.Environ.StartInstance(m1.Id(), s.StateInfo(c), nil)
-        c.Assert(err, IsNil)
-        err = m1.SetInstanceId(inst1.Id())
-        c.Assert(err, IsNil)
-
-        m2, err := s.State.AddMachine()
-        c.Assert(err, IsNil)
-        inst2, err := s.Conn.Environ.StartInstance(m2.Id(), s.StateInfo(c), nil)
-        c.Assert(err, IsNil)
-        err = m2.SetInstanceId(inst2.Id())
-        c.Assert(err, IsNil)
-
         svc, err := s.State.AddService("wordpress", s.charm)
         c.Assert(err, IsNil)
         err = svc.SetExposed()
         c.Assert(err, IsNil)
-        u1, err := svc.AddUnit()
-        c.Assert(err, IsNil)
-        err = u1.AssignToMachine(m1)
-        c.Assert(err, IsNil)
+
+        u1, m1 := s.addUnit(c, svc)
+        inst1 := s.startInstance(c, m1)
         err = u1.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
-        u2, err := svc.AddUnit()
-        c.Assert(err, IsNil)
-        err = u2.AssignToMachine(m2)
-        c.Assert(err, IsNil)
+
+        u2, m2 := s.addUnit(c, svc)
+        inst2 := s.startInstance(c, m2)
         err = u2.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst1, m1.Id(), []state.Port{{"tcp", 80}})
         assertPorts(c, inst2, m2.Id(), []state.Port{{"tcp", 80}})
 
         err = u1.ClosePort("tcp", 80)
         c.Assert(err, IsNil)
         err = u2.ClosePort("tcp", 80)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst1, m1.Id(), nil)
         assertPorts(c, inst2, m2.Id(), nil)
 }
 
 func (s *FirewallerSuite) TestFirewallerStartWithState(c *C) {
-        m, err := s.State.AddMachine()
-        c.Assert(err, IsNil)
-        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)
-
         svc, err := s.State.AddService("wordpress", s.charm)
         c.Assert(err, IsNil)
         err = svc.SetExposed()
         c.Assert(err, IsNil)
-        u, err := svc.AddUnit()
-        c.Assert(err, IsNil)
-        err = u.AssignToMachine(m)
-        c.Assert(err, IsNil)
+        u, m := s.addUnit(c, svc)
+        inst := s.startInstance(c, m)
+
         err = u.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
         err = u.OpenPort("tcp", 8080)
         c.Assert(err, IsNil)
 
         // Nothing open without firewaller.
         assertPorts(c, inst, m.Id(), nil)
 
         // Starting the firewaller opens the ports.
         fw := firewaller.NewFirewaller(s.State)
         defer func() { c.Assert(fw.Stop(), IsNil) }()
 
         assertPorts(c, inst, m.Id(), []state.Port{{"tcp", 80}, {"tcp", 8080}})
+
+        err = svc.SetExposed()
+        c.Assert(err, IsNil)
 }
 
 func (s *FirewallerSuite) TestFirewallerStartWithPartialState(c *C) {
         m, err := s.State.AddMachine()
         c.Assert(err, IsNil)
         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 16 matching lines @@
         err = u.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), []state.Port{{"tcp", 80}})
 }
 
 func (s *FirewallerSuite) TestSetClearExposedService(c *C) {
         fw := firewaller.NewFirewaller(s.State)
         defer func() { c.Assert(fw.Stop(), IsNil) }()
 
-        m, err := s.State.AddMachine()
-        c.Assert(err, IsNil)
-        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)
         svc, err := s.State.AddService("wordpress", s.charm)
         c.Assert(err, IsNil)
-        u, err := svc.AddUnit()
-        c.Assert(err, IsNil)
-        err = u.AssignToMachine(m)
-        c.Assert(err, IsNil)
+
+        u, m := s.addUnit(c, svc)
+        inst := s.startInstance(c, m)
         err = u.OpenPort("tcp", 80)
         c.Assert(err, IsNil)
         err = u.OpenPort("tcp", 8080)
         c.Assert(err, IsNil)
 
         // Not exposed service, so no open port.
         assertPorts(c, inst, m.Id(), nil)
 
         // SeExposed opens the ports.
         err = svc.SetExposed()
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), []state.Port{{"tcp", 80}, {"tcp", 8080}})
 
         // ClearExposed closes the ports again.
         err = svc.ClearExposed()
         c.Assert(err, IsNil)
 
         assertPorts(c, inst, m.Id(), nil)
 }
 
 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")
 }