added training test

hotosm · Apr 2, 2024 · af30aca · af30aca
1 parent e1cb21a
commit af30aca
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/backend/tests/test_core_model.py b/backend/tests/test_core_model.py
@@ -6,6 +6,7 @@
     AoiFactory,
     LabelFactory,
     ModelFactory,
+    TrainingFactory,
 )
 
 
@@ -17,6 +18,7 @@ def setUp(self):
         self.aoi = AoiFactory(dataset=self.dataset)
         self.label = LabelFactory(aoi=self.aoi)
         self.model = ModelFactory(dataset=self.dataset, created_by=self.user)
+        self.training = TrainingFactory(model=self.model, created_by=self.user)
 
     def test_dataset_creation(self):
         self.assertEqual(self.dataset.name, "Test Dataset")
@@ -34,3 +36,11 @@ def test_model_creation(self):
         self.assertEqual(self.model.dataset, self.dataset)
         self.assertEqual(self.model.created_by, self.user)
         self.assertEqual(self.model.status, -1)
+
+    def test_training_creation(self):
+        self.assertEqual(self.training.model, self.model)
+        self.assertEqual(self.training.status, "SUBMITTED")
+        self.assertEqual(self.training.zoom_level, [20, 21])
+        self.assertEqual(self.training.created_by, self.user)
+        self.assertEqual(self.training.epochs, 3)
+        self.assertEqual(self.training.batch_size, 24)