diff --git a/daemon/handlers.track_test.go b/daemon/handlers.track_test.go index 83ff58a..4e8a80e 100644 --- a/daemon/handlers.track_test.go +++ b/daemon/handlers.track_test.go @@ -2,6 +2,7 @@ package daemon import ( "context" + "encoding/json" "net/http" "net/http/httptest" "testing" @@ -23,6 +24,8 @@ type fakeCommandStore struct { cursorSetCalls int pruneCalls int + + engine string } func (f *fakeCommandStore) SavePre(ctx context.Context, cmd model.Command, rt time.Time) error { @@ -72,6 +75,13 @@ func (f *fakeCommandStore) Prune(ctx context.Context, cursor time.Time) error { return nil } +func (f *fakeCommandStore) Engine() string { + if f.engine == "" { + return model.StorageEngineFile + } + return f.engine +} + func (f *fakeCommandStore) Close() error { return nil } // fakeConfigService implements model.ConfigService without mockery. @@ -184,12 +194,14 @@ func (s *TrackHandlerTestSuite) TestTrackPostNotEnoughToFlush() { } func (s *TrackHandlerTestSuite) TestTrackPostFlushSyncsAndPrunes() { + var sentPayload model.PostTrackArgs server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewDecoder(r.Body).Decode(&sentPayload) w.WriteHeader(http.StatusNoContent) })) defer server.Close() - store := &fakeCommandStore{noCursorExist: true} + store := &fakeCommandStore{noCursorExist: true, engine: model.StorageEngineBolt} commandStore = store stConfig = fakeConfigService{cfg: model.ShellTimeConfig{Token: "t", APIEndpoint: server.URL, FlushCount: 1}} @@ -204,6 +216,9 @@ func (s *TrackHandlerTestSuite) TestTrackPostFlushSyncsAndPrunes() { assert.Len(s.T(), store.post, 1) assert.Equal(s.T(), 1, store.cursorSetCalls) assert.Equal(s.T(), 1, store.pruneCalls) + // the payload must report the engine that buffered the commands + assert.Equal(s.T(), model.StorageEngineBolt, sentPayload.Meta.CliEngine) + assert.Equal(s.T(), 1, sentPayload.Meta.Source, "daemon path must mark source as daemon") } func (s *TrackHandlerTestSuite) TestTrackPostFallsBackToFileStore() { diff --git a/model/api.go b/model/api.go index 27ceaaa..bd8bd55 100644 --- a/model/api.go +++ b/model/api.go @@ -32,6 +32,10 @@ type TrackingMetaData struct { Terminal string `json:"terminal,omitempty"` Multiplexer string `json:"multiplexer,omitempty"` + // CliEngine is the storage engine that buffered these commands: + // StorageEngineFile or StorageEngineBolt. + CliEngine string `json:"cliEngine,omitempty"` + // 0: cli, 1: daemon Source int `json:"source"` } diff --git a/model/store.go b/model/store.go index 33ea0f4..db11675 100644 --- a/model/store.go +++ b/model/store.go @@ -42,6 +42,11 @@ type CommandStore interface { // before the cursor), keeping unfinished pre commands. Prune(ctx context.Context, cursor time.Time) error + // Engine reports which storage engine backs this store (StorageEngineFile + // or StorageEngineBolt). It is attached to sync metadata so the server + // knows how the commands were buffered. + Engine() string + // Close releases any resources held by the store (no-op for fileStore). Close() error } diff --git a/model/store_bolt.go b/model/store_bolt.go index ecdbf64..454319b 100644 --- a/model/store_bolt.go +++ b/model/store_bolt.go @@ -241,6 +241,8 @@ func (s *boltStore) Prune(ctx context.Context, cursor time.Time) error { }) } +func (s *boltStore) Engine() string { return StorageEngineBolt } + func (s *boltStore) Close() error { if s.db == nil { return nil diff --git a/model/store_file.go b/model/store_file.go index 5e7bbab..7c9f03e 100644 --- a/model/store_file.go +++ b/model/store_file.go @@ -155,4 +155,6 @@ func (s *fileStore) Prune(ctx context.Context, cursor time.Time) error { return os.WriteFile(GetCursorFilePath(), []byte(fmt.Sprintf("%d", cursor.UnixNano())), 0644) } +func (s *fileStore) Engine() string { return StorageEngineFile } + func (s *fileStore) Close() error { return nil } diff --git a/model/tracking_build.go b/model/tracking_build.go index 859be79..204f883 100644 --- a/model/tracking_build.go +++ b/model/tracking_build.go @@ -56,6 +56,7 @@ func BuildTrackingData(ctx context.Context, store CommandStore, config ShellTime meta := TrackingMetaData{ OS: sysInfo.Os, OSVersion: sysInfo.Version, + CliEngine: store.Engine(), } trackingData := make([]TrackingData, 0) diff --git a/model/tracking_build_test.go b/model/tracking_build_test.go index 3a04102..7e12c00 100644 --- a/model/tracking_build_test.go +++ b/model/tracking_build_test.go @@ -32,6 +32,28 @@ func TestBuildTrackingData(t *testing.T) { require.Equal(t, 42, res.Data[0].PPID) require.Equal(t, "h", res.Meta.Hostname) require.Equal(t, "bash", res.Meta.Shell) + require.Equal(t, StorageEngineBolt, res.Meta.CliEngine) +} + +func TestBuildTrackingDataFileEngine(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + InitFolder("") // reset globals to the default .shelltime under the temp HOME + + store := NewFileStore() + defer store.Close() + + ctx := context.Background() + start := time.Now() + cmd := Command{Shell: "zsh", SessionID: 3, Command: "ls -la", Username: "u", Hostname: "h", Time: start} + require.NoError(t, store.SavePre(ctx, cmd, start)) + post := cmd + post.Time = start.Add(time.Second) + require.NoError(t, store.SavePost(ctx, post, 0, post.Time)) + + res, err := BuildTrackingData(ctx, store, ShellTimeConfig{}) + require.NoError(t, err) + require.Len(t, res.Data, 1) + require.Equal(t, StorageEngineFile, res.Meta.CliEngine) } func TestBuildTrackingDataExcludes(t *testing.T) {