Add test fixture

Author: vanessavmac

ami/ml/models/__init__.py

@@ -1,6 +1,6 @@
-from .algorithm import Algorithm
-from .backend import Backend
-from .pipeline import Pipeline
+from ami.ml.models.algorithm import Algorithm
+from ami.ml.models.backend import Backend
+from ami.ml.models.pipeline import Pipeline
 
 __all__ = [
     "Algorithm",

ami/ml/models/backend.py

@@ -8,11 +8,10 @@
 
 from ami.base.models import BaseModel
 from ami.main.models import Project
+from ami.ml.models.algorithm import Algorithm
+from ami.ml.models.pipeline import Pipeline
 from ami.ml.schemas import BackendStatusResponse, PipelineRegistrationResponse
 
-from .algorithm import Algorithm
-from .pipeline import Pipeline
-
 logger = logging.getLogger(__name__)
 
 
@@ -22,7 +21,7 @@ class Backend(BaseModel):
 
     projects = models.ManyToManyField("main.Project", related_name="backends", blank=True)
     endpoint_url = models.CharField(max_length=1024, null=True, blank=True)
-    pipelines = models.ManyToManyField(Pipeline, related_name="backends", blank=True)
+    pipelines = models.ManyToManyField("ml.Pipeline", related_name="backends", blank=True)
 
     def __str__(self):
         return self.endpoint_url

ami/ml/models/pipeline.py

@@ -30,11 +30,10 @@
     TaxonRank,
     update_calculated_fields_for_events,
 )
+from ami.ml.models.algorithm import Algorithm
 from ami.ml.schemas import PipelineRequest, PipelineResponse, SourceImageRequest
 from ami.ml.tasks import celery_app, create_detection_images
 
-from .algorithm import Algorithm
-
 logger = logging.getLogger(__name__)
 
 
@@ -414,7 +413,7 @@ class Pipeline(BaseModel):
     version = models.IntegerField(default=1)
     version_name = models.CharField(max_length=255, blank=True)
     # @TODO the algorithms list be retrieved by querying the pipeline endpoint
-    algorithms = models.ManyToManyField(Algorithm, related_name="pipelines")
+    algorithms = models.ManyToManyField("ml.Algorithm", related_name="pipelines")
     stages: list[PipelineStage] = SchemaField(
         default=default_stages,
         help_text=(

ami/ml/tests.py

@@ -13,19 +13,24 @@
     PipelineResponse,
     SourceImageResponse,
 )
-from ami.tests.fixtures.main import create_captures_from_files, create_ml_pipeline, setup_test_project
+from ami.tests.fixtures.main import create_captures_from_files, create_ml_backends, setup_test_project
 
 
 class TestPipelineWithMLBackend(TestCase):
     def setUp(self):
         self.project, self.deployment = setup_test_project()
         self.captures = create_captures_from_files(self.deployment, skip_existing=False)
         self.test_images = [image for image, frame in self.captures]
-        self.pipeline = create_ml_pipeline(self.project)
+        self.backend_instance = create_ml_backends(self.project)
+        self.backend = self.backend_instance
+        # @TODO: Create function to get most recent OK backend
+        self.pipeline = self.backend_instance.pipelines.all().filter(slug="constant").first()
+        self.backend_id = self.pipeline.backends.first().pk
+        # @TODO: Add error or info messages to the response if image already processed or no detections returned
 
     def test_run_pipeline(self):
         # Send images to ML backend to process and return detections
-        pipeline_response = self.pipeline.process_images(self.test_images)
+        pipeline_response = self.pipeline.process_images(self.test_images, backend_id=self.backend_id, job_id=None)
         assert pipeline_response.detections
 
 

ami/tests/fixtures/main.py

@@ -17,6 +17,7 @@
     TaxonRank,
     group_images_into_events,
 )
+from ami.ml.models.backend import Backend
 from ami.ml.tasks import create_detection_images
 from ami.tests.fixtures.storage import GeneratedTestFrame, create_storage_source, populate_bucket
 
@@ -33,62 +34,56 @@ def update_site_settings(**kwargs):
     return site
 
 
-def create_ml_pipeline(project):
-    from ami.ml.models import Algorithm, Pipeline
-
-    pipelines_to_add = [
+# @TODO: To test this: delete project in admin, then run migrate
+# (this will execute the signal in ami-platform/ami/tests/fixtures/signals.py)
+def create_ml_backends(project):
+    backends_to_add = [
         {
-            "name": "ML Dummy Backend",
-            "slug": "dummy",
-            "version": 1,
-            "algorithms": [
-                {"name": "Dummy Detector", "key": 1},
-                {"name": "Random Detector", "key": 2},
-                {"name": "Always Moth Classifier", "key": 3},
-            ],
-            "projects": {"name": project.name},
-            "endpoint_url": "http://ml_backend:2000/pipeline/process",
+            "projects": [{"name": project.name}],
+            "endpoint_url": "http://ml_backend:2000",
         },
     ]
 
-    for pipeline_data in pipelines_to_add:
-        pipeline, created = Pipeline.objects.get_or_create(
-            name=pipeline_data["name"],
-            slug=pipeline_data["slug"],
-            version=pipeline_data["version"],
-            endpoint_url=pipeline_data["endpoint_url"],
+    for backend_data in backends_to_add:
+        backend, created = Backend.objects.get_or_create(
+            endpoint_url=backend_data["endpoint_url"],
         )
 
         if created:
-            logger.info(f'Successfully created {pipeline_data["name"]}.')
+            logger.info(f'Successfully created backend with {backend_data["endpoint_url"]}.')
         else:
-            logger.info(f'Using existing pipeline {pipeline_data["name"]}.')
+            logger.info(f'Using existing backend with {backend_data["endpoint_url"]}.')
+
+        for project_data in backend_data["projects"]:
+            try:
+                project = Project.objects.get(name=project_data["name"])
+                backend.projects.add(project)
+            except Exception:
+                logger.error(f'Could not find project {project_data["name"]}.')
 
-        for algorithm_data in pipeline_data["algorithms"]:
-            algorithm, _ = Algorithm.objects.get_or_create(name=algorithm_data["name"], key=algorithm_data["key"])
-            pipeline.algorithms.add(algorithm)
+        backend.save()
 
-        pipeline.save()
+    backend.create_pipelines()
 
-    return pipeline
+    return backend
 
 
 def setup_test_project(reuse=True) -> tuple[Project, Deployment]:
+    short_id = "1ed10463"
     if reuse:
-        project, _ = Project.objects.get_or_create(name="Test Project")
+        project, _ = Project.objects.get_or_create(name=f"Test Project {short_id}")
         data_source = create_storage_source(project, "Test Data Source")
         deployment, _ = Deployment.objects.get_or_create(
             project=project, name="Test Deployment", defaults=dict(data_source=data_source)
         )
-        create_ml_pipeline(project)
+        create_ml_backends(project)
     else:
-        short_id = uuid.uuid4().hex[:8]
         project = Project.objects.create(name=f"Test Project {short_id}")
         data_source = create_storage_source(project, f"Test Data Source {short_id}")
         deployment = Deployment.objects.create(
             project=project, name=f"Test Deployment {short_id}", data_source=data_source
         )
-        create_ml_pipeline(project)
+        create_ml_backends(project)
     return project, deployment