Adding the changed ingest_companion_relations function and tests for …

…companion relations.

simple/utils/companions.py

@@ -1,7 +1,7 @@
 import logging
 import sqlalchemy.exc
 from astrodb_utils import (
-    AstroDBError,
+    AstroDBError, find_source_in_db
 )
 
 
@@ -90,6 +90,14 @@ def ingest_companion_relationships(
         msg = f"{source}: Source cannot be the same as companion name"
         logger.error(msg)
         raise AstroDBError(msg)
+
+    source_name = find_source_in_db(db, source)
+    if len(source_name) != 1:
+        msg = f"{source}: No source or multiple sources found: {source_name}"
+        logger.error(msg)
+        raise AstroDBError(msg)
+    else:
+        source_name = source_name[0]
 
     if projected_separation_arcsec is not None and projected_separation_arcsec < 0:
         msg = f"Projected separation: {projected_separation_arcsec}, cannot be negative"
@@ -118,7 +126,7 @@ def ingest_companion_relationships(
             conn.execute(
                 db.CompanionRelationships.insert().values(
                     {
-                        "source": source,
+                        "source": source_name,
                         "companion_name": companion_name,
                         "projected_separation_arcsec": projected_separation_arcsec,
                         "projected_separation_error": projected_separation_error,
@@ -133,7 +141,7 @@ def ingest_companion_relationships(
         logger.info(
             "ComapnionRelationship added: ",
             [
-                source,
+                source_name,
                 companion_name,
                 relationship,
                 projected_separation_arcsec,

tests/test_data.py

@@ -627,3 +627,15 @@ def test_radial_velocities(db):
         .astropy()
     )
     assert len(t) == 89, f"found {len(t)} radial velociies with no uncertainty"
+
+def test_companion_relations(db):
+    t = db.query(db.CompanionRelationships).astropy()
+    assert len(t) == 102, f"found {len(t)} companion relationships"
+
+    ref = "Roth24"
+    t = (
+        db.query(db.CompanionRelationships)
+        .filter(db.CompanionRelationships.c.reference == ref)
+        .astropy()
+    )
+    assert len(t) == 89, f"found {len(t)} companion relationships with {ref} reference"