Implement key management methods in cliSession with enhanced test coverage - #33
Conversation
…age for key operations
WalkthroughAdds a public ICliSession interface and implements SetKey/GetKey/DeleteKey with chunked key/value storage and legacy fallback; refactors GetRawToken/SetRawToken to use those methods; updates NewCliSession return type; expands CLI session tests and propagates SetRawToken errors from ExchangeCode. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI as CLI Client
participant Sess as CliSession (ICliSession)
participant KR as Keyring
rect rgb(240,248,255)
Note over CLI,Sess: Store token (SetRawToken)
CLI->>Sess: SetRawToken(token)
alt token == nil
Sess->>Sess: DeleteKey(baseKey)
Sess->>KR: Remove(baseKey, chunk_*, chunk_count)
KR-->>Sess: ok / not found / error
else token != nil
Sess->>Sess: Marshal(token) -> SetKey(baseKey, bytes)
Sess->>KR: Remove(existing chunks)
Sess->>KR: Set(baseKey_chunk_0..N) + Set(chunk_count)
KR-->>Sess: ok / error
end
Sess-->>CLI: success / error
end
sequenceDiagram
participant CLI as CLI Client
participant Sess as CliSession (ICliSession)
participant KR as Keyring
rect rgb(245,255,250)
Note over CLI,Sess: Retrieve token (GetRawToken)
CLI->>Sess: GetRawToken()
Sess->>Sess: GetKey(baseKey)
alt chunk_count exists
Sess->>KR: Get(chunk_0..N)
KR-->>Sess: bytes / error
Sess->>Sess: concat bytes -> unmarshal token
else legacy single key
Sess->>KR: Get(baseKey)
KR-->>Sess: bytes / not found / error
Sess->>Sess: unmarshal / return not found
end
Sess-->>CLI: token / error
end
sequenceDiagram
participant Client as OAuth2 Client
participant Flow as ExchangeCode
participant Server as Auth Server
participant Sess as sessionHooks
Client->>Flow: ExchangeCode(authCode)
Flow->>Server: token request
Server-->>Flow: token / error
alt token received
Flow->>Sess: SetRawToken(token)
Sess-->>Flow: ok / error
alt SetRawToken error
Flow-->>Client: return error
else
Flow-->>Client: success
end
else error
Flow-->>Client: return error
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Tip For best results, initiate chat on the files or code changes. @evgenyk Hello! I see you've mentioned me. How can I help you with this PR? I'm ready to assist with code review, analysis, or any questions you might have about the key management methods implementation and test coverage. |
|
@CodeRabbit code review please |
|
@evgenyk I'll review the PR changes for you right now, focusing on the key management methods implementation in cliSession and the test coverage enhancements. ✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
frameworks/cli/cliSession.go (1)
168-213: SetKey implementation needs attention to error handling and cleanupThere are several issues with the current implementation:
Variable shadowing (Line 180): The
chunksvariable from Line 177 is shadowed by the loop variable, which could cause confusion.Silent failure on chunk removal (Line 184-186): Similar to
DeleteKey, breaking on error leaves partial chunks in place without proper error reporting.Duplicate countKey assignment (Line 204): The
countKeyvariable is reassigned with the same value it already has from Line 169.Here's an improved implementation:
func (c *cliSession) SetKey(key string, value []byte) error { countKey := fmt.Sprintf("%s_chunk_count", key) if len(value) == 0 { // Remove legacy single-key entry return c.DeleteKey(key) } const chunkSize = 1024 - chunks := (len(value) + chunkSize - 1) / chunkSize + numChunks := (len(value) + chunkSize - 1) / chunkSize // Try to get chunk count - if chunks, err := c.getChunkCount(countKey); err == nil { + if existingChunks, err := c.getChunkCount(countKey); err == nil { // Remove any existing chunks - for i := range chunks { + for i := range existingChunks { chunkKey := fmt.Sprintf("%s_chunk_%d", key, i) if err := c.keyring.Remove(chunkKey); err != nil { - break + if !errors.Is(err, keyring.ErrKeyNotFound) { + // Best effort cleanup - continue removing other chunks + // but return error at the end + } } } } // Save chunks - for i := range chunks { + for i := range numChunks { start := i * chunkSize end := min(start+chunkSize, len(value)) chunkKey := fmt.Sprintf("%s_chunk_%d", key, i) if err := c.keyring.Set(keyring.Item{ Key: chunkKey, Data: value[start:end], }); err != nil { return fmt.Errorf("failed to save token chunk %d: %w", i, err) } } // Save chunk count - countKey = fmt.Sprintf("%s_chunk_count", key) - countData := fmt.Appendf(nil, "%d", chunks) + countData := fmt.Appendf(nil, "%d", numChunks) if err := c.keyring.Set(keyring.Item{ Key: countKey, Data: countData, }); err != nil { return fmt.Errorf("failed to save chunk count: %w", err) } return nil }
🧹 Nitpick comments (2)
frameworks/cli/cliSession_test.go (1)
296-321: Consider enhancing the chunk removal error testThe test correctly verifies that chunk removal errors don't propagate (Line 316), but the comment on Line 320 indicates that
chunk_1is "not removed due to break". However, looking at theDeleteKeyimplementation incliSession.go, the loop does break on error (Line 53), which would leave subsequent chunks intact.Consider adding an assertion to verify that the chunk count key is also not removed when chunk deletion fails, as this could lead to inconsistent state.
func TestCliSession_DeleteKey_ChunkRemoveError(t *testing.T) { assert := assert.New(t) mk := &mockKeyring{ items: map[string]keyring.Item{ "mykey_chunk_count": {Key: "mykey_chunk_count", Data: []byte("2")}, "mykey_chunk_0": {Key: "mykey_chunk_0", Data: []byte("chunk1")}, "mykey_chunk_1": {Key: "mykey_chunk_1", Data: []byte("chunk2")}, }, getErrs: map[string]error{}, } // Simulate error on removing chunk_0 mk.RemoveFunc = func(key string) error { if key == "mykey_chunk_0" { return errors.New("remove error") } delete(mk.items, key) return nil } cs := &cliSession{keyring: mk} err := cs.DeleteKey("mykey") assert.Nil(err) // Should not propagate error _, exists0 := mk.items["mykey_chunk_0"] _, exists1 := mk.items["mykey_chunk_1"] + _, existsCount := mk.items["mykey_chunk_count"] assert.True(exists0) // chunk_0 not removed due to error assert.True(exists1) // chunk_1 not removed due to break + assert.True(existsCount) // chunk_count should also remain due to early return }frameworks/cli/cliSession.go (1)
160-166: Redundant nil return statementThe explicit
return nilon Line 165 is redundant sinceerrwill already be nil ifSetKeysucceeds. Consider simplifying:err = c.SetKey(key, t) - if err != nil { - return err - } - - return nil + return err
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
frameworks/cli/cliSession.go(6 hunks)frameworks/cli/cliSession_test.go(3 hunks)oauth2/authorization_code/authorization_code.go(1 hunks)
🔇 Additional comments (9)
oauth2/authorization_code/authorization_code.go (1)
261-263: Good fix for proper error propagationThe change correctly propagates the error from
SetRawToken, ensuring that token storage failures are properly handled and reported to the caller. This is an important improvement for error handling.frameworks/cli/cliSession_test.go (3)
27-30: LGTM! Flexible mock design for testing error pathsThe addition of
RemoveFuncto the mock allows for simulating removal errors in tests, which is essential for comprehensive error path coverage.
42-54: Well-structured mock methods with proper error injection supportThe refactored mock methods maintain clear semantics while enabling controllable error simulation through
RemoveFunc. The implementation correctly handles the delete operation whenRemoveFuncis not set.
162-332: Excellent test coverage for key management operationsThe test suite provides comprehensive coverage for:
- Single and chunked key operations
- Error handling scenarios (parse errors, missing chunks, removal failures)
- Edge cases (non-existent keys, empty operations)
- Legacy key cleanup during chunked operations
The tests are well-structured and thoroughly validate the chunked storage implementation.
frameworks/cli/cliSession.go (5)
26-31: Good interface design for extended key managementThe
ICliSessioninterface cleanly extendsISessionHookswith generic key management methods, providing a clear separation of concerns and improved API surface for CLI session operations.
67-92: LGTM! Clean implementation of GetKey with backwards compatibilityThe method correctly handles both chunked and single-key storage formats, providing seamless backwards compatibility for existing data.
135-138: Improved error message clarityThe error message has been updated from "failed to get token" to "failed to read token", which better describes the operation being performed.
152-153: Good simplification using DeleteKeyUsing
DeleteKeyfor nil token handling ensures consistent cleanup of both single and chunked token data.
244-244: Good API evolution with backwards compatibilityThe return type change from
ISessionHookstoICliSessionmaintains backwards compatibility while exposing the new key management functionality.
…ontinued processing
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
frameworks/cli/cliSession_test.go (3)
49-54: Guard against nil map in Set and return actual keys from KeysMakes the mock safer and more generally useful.
Apply this diff:
func (m *mockKeyring) Set(item keyring.Item) error { - m.items[item.Key] = item + if m.items == nil { + m.items = map[string]keyring.Item{} + } + m.items[item.Key] = item return nil } -func (m *mockKeyring) Keys() ([]string, error) { return nil, nil } +func (m *mockKeyring) Keys() ([]string, error) { + keys := make([]string, 0, len(m.items)) + for k := range m.items { + keys = append(keys, k) + } + return keys, nil +} func (m *mockKeyring) Reset() error { m.items = map[string]keyring.Item{}; return nil }
208-221: Use “key” not “token” in GetKey error assertionsThese tests target GetKey; aligning messages avoids confusion with token-specific flows.
Apply this diff:
- assert.Contains(err.Error(), "failed to get token") + assert.Contains(err.Error(), "failed to get key")- assert.Contains(err.Error(), "failed to get token chunk 0") + assert.Contains(err.Error(), "failed to get key chunk 0")- assert.Contains(err.Error(), "failed to get token") + assert.Contains(err.Error(), "failed to get key")If the implementation intentionally reuses token-oriented messages for generic keys, please confirm and we can keep the current assertions.
Also applies to: 223-238, 240-254
296-323: Consider asserting chunk-count key behavior on partial failureWhen chunk_0 removal fails, decide whether chunk_count should be retained or also removed and assert accordingly to lock semantics. Optionally add a test for failure removing the chunk_count key itself.
Do you want me to draft a companion test (e.g., DeleteKey_ChunkCountRemoveError) that simulates an error on removing "mykey_chunk_count"?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
frameworks/cli/cliSession.go(6 hunks)frameworks/cli/cliSession_test.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- frameworks/cli/cliSession.go
🔇 Additional comments (7)
frameworks/cli/cliSession_test.go (7)
27-30: Mock extension for removal error injection looks goodAdding RemoveFunc provides precise control in tests without complicating the happy path.
42-48: Remove delegation + no-op on missing keys — LGTMBehavior is predictable and fits the tests’ needs.
161-171: DeleteKey should be no-op when key is absent — correctThe test asserts the desired behavior.
172-187: Single-key retrieval path covered wellHappy-path GetKey without chunking is validated.
189-206: Chunked retrieval path covered wellTwo-chunk reconstruction is exercised cleanly.
255-270: Single-key delete path validatedState check post-delete is clear.
272-294: Chunked delete path validated, including legacy fallbackVerifies all fragments are removed — nice coverage.
1 participant
Introduce methods for setting, getting, and deleting keys in cliSession, along with improved test coverage for these operations. This update enhances the key management functionality and ensures robustness through comprehensive testing.