@@ -673,14 +673,26 @@ func TestAckHandleUpgrade(t *testing.T) {
673673 name string
674674 event UpgradeEvent
675675 bulker func (t * testing.T ) * ftesting.MockBulk
676+ agent * model.Agent
676677 }{{
677678 name : "ok" ,
678679 event : UpgradeEvent {},
679680 bulker : func (t * testing.T ) * ftesting.MockBulk {
680681 m := ftesting .NewMockBulk ()
681- m .On ("Update" , mock .Anything , mock .Anything , mock .Anything , mock .Anything , mock .Anything ).Return (nil ).Once ()
682+ m .On ("Update" , mock .Anything , mock .Anything , mock .Anything , mock .MatchedBy (func (p []byte ) bool {
683+ var body map [string ]map [string ]interface {}
684+ if err := json .Unmarshal (p , & body ); err != nil {
685+ t .Fatal (err )
686+ }
687+ upgradeAttempts , ok := body ["doc" ]["upgrade_attempts" ]
688+ return ok == true && upgradeAttempts == nil
689+ }), mock .Anything ).Return (nil ).Once ()
682690 return m
683691 },
692+ agent : & model.Agent {
693+ ESDocument : model.ESDocument {Id : "ab12dcd8-bde0-4045-92dc-c4b27668d735" },
694+ Agent : & model.AgentMetadata {Version : "8.0.0" },
695+ },
684696 }, {
685697 name : "retry signaled" ,
686698 event : UpgradeEvent {
@@ -708,6 +720,10 @@ func TestAckHandleUpgrade(t *testing.T) {
708720 }), mock .Anything ).Return (nil ).Once ()
709721 return m
710722 },
723+ agent : & model.Agent {
724+ ESDocument : model.ESDocument {Id : "ab12dcd8-bde0-4045-92dc-c4b27668d735" },
725+ Agent : & model.AgentMetadata {Version : "8.0.0" },
726+ },
711727 }, {
712728 name : "no more retries" ,
713729 event : UpgradeEvent {
@@ -734,14 +750,34 @@ func TestAckHandleUpgrade(t *testing.T) {
734750 }), mock .Anything ).Return (nil ).Once ()
735751 return m
736752 },
753+ agent : & model.Agent {
754+ ESDocument : model.ESDocument {Id : "ab12dcd8-bde0-4045-92dc-c4b27668d735" },
755+ Agent : & model.AgentMetadata {Version : "8.0.0" },
756+ },
757+ }, {
758+ name : "keep upgrade_attempts if upgrade_details is not nil" ,
759+ event : UpgradeEvent {},
760+ bulker : func (t * testing.T ) * ftesting.MockBulk {
761+ m := ftesting .NewMockBulk ()
762+ m .On ("Update" , mock .Anything , mock .Anything , mock .Anything , mock .MatchedBy (func (p []byte ) bool {
763+ var body map [string ]map [string ]interface {}
764+ if err := json .Unmarshal (p , & body ); err != nil {
765+ t .Fatal (err )
766+ }
767+ _ , ok := body ["doc" ]["upgrade_attempts" ]
768+ return ok == false
769+ }), mock .Anything ).Return (nil ).Once ()
770+ return m
771+ },
772+ agent : & model.Agent {
773+ ESDocument : model.ESDocument {Id : "ab12dcd8-bde0-4045-92dc-c4b27668d735" },
774+ Agent : & model.AgentMetadata {Version : "8.0.0" },
775+ UpgradeDetails : & model.UpgradeDetails {},
776+ },
737777 }}
738778 cfg := & config.Server {
739779 Limits : config.ServerLimits {},
740780 }
741- agent := & model.Agent {
742- ESDocument : model.ESDocument {Id : "ab12dcd8-bde0-4045-92dc-c4b27668d735" },
743- Agent : & model.AgentMetadata {Version : "8.0.0" },
744- }
745781 ctx := context .Background ()
746782 cache , err := cache .New (config.Cache {NumCounters : 100 , MaxCost : 100000 })
747783 if err != nil {
@@ -753,7 +789,7 @@ func TestAckHandleUpgrade(t *testing.T) {
753789 bulker := tc .bulker (t )
754790 ack := NewAckT (cfg , bulker , cache )
755791
756- err := ack .handleUpgrade (ctx , logger , agent , tc .event )
792+ err := ack .handleUpgrade (ctx , logger , tc . agent , tc .event )
757793 assert .NoError (t , err )
758794 bulker .AssertExpectations (t )
759795 })
0 commit comments