generated from TBD54566975/tbd-project-template
-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: allow leases in module tests (#1562)
- split up the code in `ftl/leases.go` between: - the business logic, which drives the acquisition and heartbeating of the lease - the communication code, now called `leaseClient`, which streams requests and responses with the controller - `ftl/leases.go` chooses which client it uses by asking the ModuleContext for a mock lease client, and uses the standard one if no mock is provided. - `ftltest.MockLeaseClient` implements the `LeaseClient` protocol. `ftltest` provides this mock to the `modulecontext` when building the context It's odd that the `LeaseClient` protocol is in `modulecontext`, LMK if you have a better way of arranging things without making import cycles...
- Loading branch information
Showing
10 changed files
with
281 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
package ftltest | ||
|
||
import ( | ||
"context" | ||
"fmt" | ||
"strings" | ||
"sync" | ||
"time" | ||
|
||
"github.com/TBD54566975/ftl/go-runtime/ftl" | ||
"github.com/TBD54566975/ftl/internal/modulecontext" | ||
) | ||
|
||
// fakeLeaseClient is a simple in-memory lease client for testing. | ||
// | ||
// It does not include any checks on module names, as it assumes that all leases are within the module being tested | ||
type fakeLeaseClient struct { | ||
lock sync.Mutex | ||
deadlines map[string]time.Time | ||
} | ||
|
||
var _ modulecontext.LeaseClient = &fakeLeaseClient{} | ||
|
||
func newFakeLeaseClient() *fakeLeaseClient { | ||
return &fakeLeaseClient{ | ||
deadlines: make(map[string]time.Time), | ||
} | ||
} | ||
|
||
func keyForKeys(keys []string) string { | ||
return strings.Join(keys, "\n") | ||
} | ||
|
||
func (c *fakeLeaseClient) Acquire(ctx context.Context, module string, key []string, ttl time.Duration) error { | ||
c.lock.Lock() | ||
defer c.lock.Unlock() | ||
|
||
k := keyForKeys(key) | ||
if deadline, ok := c.deadlines[k]; ok { | ||
if time.Now().Before(deadline) { | ||
return ftl.ErrLeaseHeld | ||
} | ||
} | ||
|
||
c.deadlines[k] = time.Now().Add(ttl) | ||
return nil | ||
} | ||
|
||
func (c *fakeLeaseClient) Heartbeat(ctx context.Context, module string, key []string, ttl time.Duration) error { | ||
c.lock.Lock() | ||
defer c.lock.Unlock() | ||
|
||
k := keyForKeys(key) | ||
if deadline, ok := c.deadlines[k]; ok { | ||
if !time.Now().Before(deadline) { | ||
return fmt.Errorf("could not heartbeat expired lease") | ||
} | ||
c.deadlines[k] = time.Now().Add(ttl) | ||
return nil | ||
} | ||
return fmt.Errorf("could not heartbeat lease: no active lease found") | ||
} | ||
|
||
func (c *fakeLeaseClient) Release(ctx context.Context, key []string) error { | ||
c.lock.Lock() | ||
defer c.lock.Unlock() | ||
|
||
k := keyForKeys(key) | ||
if deadline, ok := c.deadlines[k]; ok { | ||
if !time.Now().Before(deadline) { | ||
return fmt.Errorf("could not release lease after timeout") | ||
} | ||
delete(c.deadlines, k) | ||
return nil | ||
} | ||
return fmt.Errorf("could not release lease: no active lease found") | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,63 @@ | ||
package ftltest | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"testing" | ||
"time" | ||
|
||
"github.com/TBD54566975/ftl/go-runtime/ftl" | ||
"github.com/alecthomas/assert/v2" | ||
) | ||
|
||
var ( | ||
keys1 = []string{"one", "1"} | ||
keys2 = []string{"two", "2"} | ||
module = "test" | ||
) | ||
|
||
func TestDoubleAcquireLease(t *testing.T) { | ||
ctx := context.Background() | ||
client := newFakeLeaseClient() | ||
|
||
// Acquire a lease, and immediately try to acquire it again. | ||
err := client.Acquire(ctx, module, keys1, 1*time.Second) | ||
assert.NoError(t, err) | ||
err = client.Acquire(ctx, module, keys1, 1*time.Second) | ||
assert.True(t, errors.Is(err, ftl.ErrLeaseHeld), "expected lease to already be held") | ||
} | ||
|
||
func TestAcquireTwoDifferentLeases(t *testing.T) { | ||
ctx := context.Background() | ||
client := newFakeLeaseClient() | ||
|
||
err := client.Acquire(ctx, module, keys1, 1*time.Second) | ||
assert.NoError(t, err) | ||
err = client.Acquire(ctx, module, keys2, 1*time.Second) | ||
assert.NoError(t, err) | ||
} | ||
|
||
func TestExpiry(t *testing.T) { | ||
ctx := context.Background() | ||
client := newFakeLeaseClient() | ||
|
||
err := client.Acquire(ctx, module, keys1, 500*time.Millisecond) | ||
assert.NoError(t, err) | ||
time.Sleep(250 * time.Millisecond) | ||
err = client.Heartbeat(ctx, module, keys1, 500*time.Millisecond) | ||
assert.NoError(t, err) | ||
time.Sleep(250 * time.Millisecond) | ||
err = client.Heartbeat(ctx, module, keys1, 500*time.Millisecond) | ||
assert.NoError(t, err) | ||
|
||
// wait longer than ttl | ||
time.Sleep(1 * time.Second) | ||
err = client.Heartbeat(ctx, module, keys1, 500*time.Millisecond) | ||
assert.Error(t, err, "expected error for heartbeating expired lease") | ||
err = client.Release(ctx, keys1) | ||
assert.Error(t, err, "expected error for heartbeating expired lease") | ||
|
||
// try and acquire again | ||
err = client.Acquire(ctx, module, keys1, 1*time.Second) | ||
assert.NoError(t, err) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,28 @@ | ||
package leases | ||
|
||
import ( | ||
"testing" | ||
"time" | ||
|
||
"github.com/TBD54566975/ftl/go-runtime/ftl/ftltest" | ||
"github.com/alecthomas/assert/v2" | ||
"golang.org/x/sync/errgroup" | ||
) | ||
|
||
func TestLease(t *testing.T) { | ||
ctx := ftltest.Context( | ||
ftltest.WithProjectFiles(), | ||
) | ||
// test that we can acquire a lease in a test environment | ||
wg := errgroup.Group{} | ||
wg.Go(func() error { | ||
return Acquire(ctx) | ||
}) | ||
|
||
// test that we get an error acquiring another lease at the same time | ||
time.Sleep(1 * time.Second) | ||
err := Acquire(ctx) | ||
assert.Error(t, err, "expected error for acquiring lease while another is held") | ||
|
||
assert.NoError(t, wg.Wait(), "expected no error acquiring the initial lease") | ||
} |
Oops, something went wrong.