github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/server/sertesting/time.go (about)

     1  // Copyright 2017 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package sertesting contains utilities useful for testing the fleetspeak server
    16  // and server components.
    17  package sertesting
    18  
    19  import (
    20  	"sync/atomic"
    21  	"time"
    22  
    23  	"github.com/google/fleetspeak/fleetspeak/src/server/internal/ftime"
    24  )
    25  
    26  // FakeTime represents a fake time which can control the time seen by the fleetspeak
    27  // system during unit tests.
    28  type FakeTime struct {
    29  	t    int64
    30  	orig func() time.Time
    31  }
    32  
    33  // Get returns the current fake time.
    34  func (t *FakeTime) Get() time.Time {
    35  	return time.Unix(atomic.LoadInt64(&t.t), 0).UTC()
    36  }
    37  
    38  // SetSeconds sets the current fake time to s seconds since epoch.
    39  func (t *FakeTime) SetSeconds(s int64) {
    40  	atomic.StoreInt64(&t.t, s)
    41  }
    42  
    43  // AddSeconds adds s seconds to the fake time.
    44  func (t *FakeTime) AddSeconds(s int64) {
    45  	atomic.AddInt64(&t.t, s)
    46  }
    47  
    48  // Revert returns the time seen by the fleetspeak system to what it was
    49  // previously.
    50  func (t *FakeTime) Revert() {
    51  	ftime.Now = t.orig
    52  }
    53  
    54  // FakeNow changes the implementation of Now used by the fleetspeak system. It
    55  // returns a FakeTime which controls the time that the system sees until Revert
    56  // is called.
    57  //
    58  // Note that this function, and FakeTime.Revert are not thread safe. In
    59  // particular it they should be called at the start and end of tests when no Fleetspeak
    60  // component is started.
    61  func FakeNow(start int64) *FakeTime {
    62  	r := FakeTime{
    63  		t:    start,
    64  		orig: ftime.Now,
    65  	}
    66  	ftime.Now = r.Get
    67  	return &r
    68  }