From 525571c32e0e8f26eac9f422ecc86d1dae82e121 Mon Sep 17 00:00:00 2001
From: hunteraraujo <hunter_araujo@msn.com>
Date: Tue, 19 Sep 2023 21:31:02 -0700
Subject: [PATCH] Enhance runBenchmark with TestSuite Tracking (#5268)

---
 frontend/lib/viewmodels/skill_tree_viewmodel.dart  | 14 ++++++++++++--
 frontend/lib/views/task_queue/task_queue_view.dart |  7 +++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/frontend/lib/viewmodels/skill_tree_viewmodel.dart b/frontend/lib/viewmodels/skill_tree_viewmodel.dart
index 70e7d61e..862efa48 100644
--- a/frontend/lib/viewmodels/skill_tree_viewmodel.dart
+++ b/frontend/lib/viewmodels/skill_tree_viewmodel.dart
@@ -5,8 +5,10 @@ import 'package:auto_gpt_flutter_client/models/skill_tree/skill_tree_edge.dart';
 import 'package:auto_gpt_flutter_client/models/skill_tree/skill_tree_node.dart';
 import 'package:auto_gpt_flutter_client/models/step.dart';
 import 'package:auto_gpt_flutter_client/models/task.dart';
+import 'package:auto_gpt_flutter_client/models/test_suite.dart';
 import 'package:auto_gpt_flutter_client/services/benchmark_service.dart';
 import 'package:auto_gpt_flutter_client/viewmodels/chat_viewmodel.dart';
+import 'package:auto_gpt_flutter_client/viewmodels/task_viewmodel.dart';
 import 'package:collection/collection.dart';
 import 'package:flutter/foundation.dart';
 import 'package:flutter/services.dart';
@@ -139,11 +141,15 @@ class SkillTreeViewModel extends ChangeNotifier {
   }
 
   // TODO: Move to task queue view model
-  // TODO: We should be creating TestSuite objects
-  Future<void> runBenchmark(ChatViewModel chatViewModel) async {
+  Future<void> runBenchmark(
+      ChatViewModel chatViewModel, TaskViewModel taskViewModel) async {
     // Clear the benchmarkStatusList
     benchmarkStatusList.clear();
 
+    // Create a new TestSuite object with the current timestamp
+    final testSuite =
+        TestSuite(timestamp: DateTime.now().toIso8601String(), tests: []);
+
     // Set the benchmark running flag to true
     isBenchmarkRunning = true;
     // Notify listeners
@@ -209,7 +215,11 @@ class SkillTreeViewModel extends ChangeNotifier {
               "Benchmark for node ${node.id} failed. Stopping all benchmarks.");
           break;
         }
+        testSuite.tests.add(task);
       }
+
+      // Add the TestSuite to the TaskViewModel
+      taskViewModel.addTestSuite(testSuite);
     } catch (e) {
       print("Error while running benchmark: $e");
     }
diff --git a/frontend/lib/views/task_queue/task_queue_view.dart b/frontend/lib/views/task_queue/task_queue_view.dart
index 3d46b8cd..1dd8ef3b 100644
--- a/frontend/lib/views/task_queue/task_queue_view.dart
+++ b/frontend/lib/views/task_queue/task_queue_view.dart
@@ -1,4 +1,5 @@
 import 'package:auto_gpt_flutter_client/viewmodels/chat_viewmodel.dart';
+import 'package:auto_gpt_flutter_client/viewmodels/task_viewmodel.dart';
 import 'package:flutter/material.dart';
 import 'package:auto_gpt_flutter_client/viewmodels/skill_tree_viewmodel.dart';
 import 'package:provider/provider.dart';
@@ -56,12 +57,14 @@ class TaskQueueView extends StatelessWidget {
                 onPressed: viewModel.isBenchmarkRunning
                     ? null
                     : () {
-                        // TODO: Handle this better
+                        // TODO: We should not be passing this dependency in like this
                         final chatViewModel =
                             Provider.of<ChatViewModel>(context, listen: false);
+                        final taskViewModel =
+                            Provider.of<TaskViewModel>(context, listen: false);
                         chatViewModel.clearCurrentTaskAndChats();
                         // Call runBenchmark method from SkillTreeViewModel
-                        viewModel.runBenchmark(chatViewModel);
+                        viewModel.runBenchmark(chatViewModel, taskViewModel);
                       },
                 child: Row(
                   mainAxisAlignment: