From 23cc356cd06ad00900518347589afabae7af9ce4 Mon Sep 17 00:00:00 2001
From: TheBurchLog <5104941+TheBurchLog@users.noreply.github.com>
Date: Fri, 7 Jun 2024 13:29:16 +0000
Subject: [PATCH] Updating Job Model

---
 brewtils/models.py        | 12 +++++++-----
 brewtils/schemas.py       |  3 ++-
 brewtils/test/fixtures.py |  5 +++--
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/brewtils/models.py b/brewtils/models.py
index 1b3f9618..ed482c72 100644
--- a/brewtils/models.py
+++ b/brewtils/models.py
@@ -1209,6 +1209,7 @@ def __init__(
         status=None,
         max_instances=None,
         timeout=None,
+        replication_id=None,
     ):
         self.id = id
         self.name = name
@@ -1225,6 +1226,7 @@ def __init__(
         self.status = status
         self.max_instances = max_instances
         self.timeout = timeout
+        self.replication_id = replication_id
 
     def __str__(self):
         return "%s: %s" % (self.name, self.id)
@@ -1729,16 +1731,16 @@ def __repr__(self):
 class Replication(BaseModel):
     schema = "ReplicationSchema"
 
-    def __init__(self, id=None, replication_id=None, heartbeat=None):
+    def __init__(self, id=None, replication_id=None, expires_at=None):
         self.id = id
         self.replication_id = replication_id
-        self.heartbeat = heartbeat
+        self.expires_at = expires_at
 
     def __str__(self):
-        return "%s:%s" % (self.replication_id, self.heartbeat)
+        return "%s:%s" % (self.replication_id, self.expires_at)
 
     def __repr__(self):
-        return "<Replication: replication_id=%s, heartbeat=%s>" % (
+        return "<Replication: replication_id=%s, expires_at=%s>" % (
             self.replication_id,
-            self.heartbeat,
+            self.expires_at,
         )
diff --git a/brewtils/schemas.py b/brewtils/schemas.py
index 1345139b..0fcc99bd 100644
--- a/brewtils/schemas.py
+++ b/brewtils/schemas.py
@@ -545,6 +545,7 @@ class JobSchema(BaseSchema):
     status = fields.Str(allow_none=True)
     max_instances = fields.Int(allow_none=True)
     timeout = fields.Int(allow_none=True)
+    replication_id = fields.Str(allow_none=True)
 
 
 class JobExportInputSchema(BaseSchema):
@@ -645,7 +646,7 @@ class TopicSchema(BaseSchema):
 class ReplicationSchema(BaseSchema):
     id = fields.Str(allow_none=True)
     replication_id = fields.Str(allow_none=True)
-    heartbeat = DateTime(allow_none=True, format="epoch", example="1500065932000")
+    expires_at = DateTime(allow_none=True, format="epoch", example="1500065932000")
 
 
 class UserSchema(BaseSchema):
diff --git a/brewtils/test/fixtures.py b/brewtils/test/fixtures.py
index 2bc43ede..195d8798 100644
--- a/brewtils/test/fixtures.py
+++ b/brewtils/test/fixtures.py
@@ -584,6 +584,7 @@ def job_dict(ts_epoch, request_template_dict, date_trigger_dict):
         "status": "RUNNING",
         "max_instances": 3,
         "timeout": 30,
+        "replication_id": "123"
     }
 
 
@@ -945,12 +946,12 @@ def replication_dict(ts_epoch):
     return {
         "id": "1234",
         "replication_id": "89cd6a3a-e0e2-486b-b8e8-535d1893faf3",
-        "heartbeat": ts_epoch,
+        "expires_at": ts_epoch,
     }
 
 
 @pytest.fixture
 def bg_replication(replication_dict, ts_dt):
     dict_copy = copy.deepcopy(replication_dict)
-    dict_copy["heartbeat"] = ts_dt
+    dict_copy["expires_at"] = ts_dt
     return Replication(**dict_copy)