Add raw data setter and tests, also improving test for the raw data property

Author: josenavas

qiita_db/metadata_template/prep_template.py

@@ -279,9 +279,30 @@ def data_type(self, ret_id=False):
     @property
     def raw_data(self):
         conn_handler = SQLConnectionHandler()
-        return conn_handler.execute_fetchone(
+        result = conn_handler.execute_fetchone(
             "SELECT raw_data_id FROM qiita.prep_template "
-            "WHERE prep_template_id=%s", (self.id,))[0]
+            "WHERE prep_template_id=%s", (self.id,))
+        if result:
+            return result[0]
+        return None
+
+    @raw_data.setter
+    def raw_data(self, raw_data):
+        conn_handler = SQLConnectionHandler()
+        sql = """SELECT (
+                    SELECT raw_data_id
+                    FROM qiita.prep_template
+                    WHERE prep_template_id=%s)
+                IS NOT NULL"""
+        exists = conn_handler.execute_fetchone(sql, (self.id,))[0]
+        if exists:
+            raise QiitaDBError(
+                "Prep template %d already has a raw data associated"
+                % self.id)
+        sql = """UPDATE qiita.prep_template
+                 SET raw_data_id = %s
+                 WHERE prep_template_id = %s"""
+        conn_handler.execute(sql, (raw_data.id, self.id))
 
     @property
     def preprocessed_data(self):

qiita_db/metadata_template/test/test_prep_template.py

@@ -24,7 +24,8 @@
                                  QiitaDBDuplicateHeaderError,
                                  QiitaDBExecutionError,
                                  QiitaDBColumnError,
-                                 QiitaDBWarning)
+                                 QiitaDBWarning,
+                                 QiitaDBError)
 from qiita_db.sql_connection import SQLConnectionHandler
 from qiita_db.study import Study
 from qiita_db.data import RawData, ProcessedData
@@ -535,10 +536,6 @@ def test_data_type_id(self):
         """data_type returns the int with the data_type_id"""
         self.assertTrue(self.tester.data_type(ret_id=True), 2)
 
-    def test_raw_data(self):
-        """Returns the raw_data associated with the prep template"""
-        self.assertEqual(self.tester.raw_data, 1)
-
     def test_preprocessed_data(self):
         """Returns the preprocessed data list generated from this template"""
         self.assertEqual(self.tester.preprocessed_data, [1, 2])
@@ -1318,6 +1315,27 @@ def test_check_restrictions(self):
              PREP_TEMPLATE_COLUMNS_TARGET_GENE['demultiplex']])
         self.assertEqual(obs, {'primer'})
 
+    def test_raw_data(self):
+        """Returns the raw_data associated with the prep template"""
+        self.assertEqual(self.tester.raw_data, 1)
+
+        pt = PrepTemplate.create(self.metadata, self.test_study,
+                                 self.data_type_id)
+        self.assertEqual(pt.raw_data, None)
+
+    def test_raw_data_setter_error(self):
+        rd = RawData(1)
+        with self.assertRaises(QiitaDBError):
+            self.tester.raw_data = rd
+
+    def test_raw_data_setter(self):
+        rd = RawData(1)
+        pt = PrepTemplate.create(self.metadata, self.test_study,
+                                 self.data_type_id)
+        self.assertEqual(pt.raw_data, None)
+        pt.raw_data = rd
+        self.assertEqual(pt.raw_data, 1)
+
 
 EXP_PREP_TEMPLATE = (
     'sample_name\tbarcode\tcenter_name\tcenter_project_name\t'