X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/f5c601f1801aaa276f8c335fcb51c373edc0b50c..bd7da3557dc7a17020931f5a49ed5007fa81a645:/daemon_message_helper_test.go?ds=inline

diff --git a/daemon_message_helper_test.go b/daemon_message_helper_test.go
index 2a627e2..6b599ae 100644
--- a/daemon_message_helper_test.go
+++ b/daemon_message_helper_test.go
@@ -36,6 +36,7 @@ func newDaemonMessageTest(t *testing.T, dir string, server *ModelFake) *DaemonMe
 	test := &DaemonMessageTest{
 		logs: &SyncBuffer{},
 	}
+
 	test.agents, err = openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, newTestClock().Now, daemonTestLogger(test.logs))
 	require.NoError(t, err)
 	t.Cleanup(func() {
@@ -57,6 +58,25 @@ func newDaemonMessageTest(t *testing.T, dir string, server *ModelFake) *DaemonMe
 	return test
 }
 
+func (r *DaemonMessageTest) Close(t *testing.T) {
+	t.Helper()
+
+	if r.cancel != nil {
+		r.cancel()
+
+		for _, done := range r.done {
+			require.ErrorIs(t, <-done, context.Canceled)
+		}
+
+		r.cancel = nil
+		r.done = nil
+	}
+
+	require.NoError(t, r.agents.Close())
+
+	r.agents = nil
+}
+
 func (r *DaemonMessageTest) Start(t *testing.T, indices ...int) {
 	t.Helper()
 
@@ -83,22 +103,3 @@ func (r *DaemonMessageTest) Queue(t *testing.T) DaemonMessageFile {
 
 	return file
 }
-
-func (r *DaemonMessageTest) Close(t *testing.T) {
-	t.Helper()
-
-	if r.cancel != nil {
-		r.cancel()
-
-		for _, done := range r.done {
-			require.ErrorIs(t, <-done, context.Canceled)
-		}
-
-		r.cancel = nil
-		r.done = nil
-	}
-
-	require.NoError(t, r.agents.Close())
-
-	r.agents = nil
-}