added training factory

hotosm · Apr 2, 2024 · e1cb21a · e1cb21a
1 parent 93965fd
commit e1cb21a
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/backend/tests/factories.py b/backend/tests/factories.py
@@ -1,7 +1,7 @@
 import factory
 
 from login.models import OsmUser
-from core.models import Dataset, AOI, Label, Model
+from core.models import Dataset, AOI, Label, Model, Training
 
 
 class OsmUserFactory(factory.django.DjangoModelFactory):
@@ -48,3 +48,16 @@ class Meta:
     name = "Test Model"
     created_by = factory.SubFactory(OsmUserFactory)
     status = -1
+
+
+class TrainingFactory(factory.django.DjangoModelFactory):
+
+    class Meta:
+        model = Training
+
+    model = factory.SubFactory(ModelFactory)
+    status = "SUBMITTED"
+    zoom_level = [20, 21]
+    created_by = factory.SubFactory(OsmUserFactory)
+    epochs = 3
+    batch_size = 24