Merge pull request #11 from juneHQ/vinayak/2024-12-02-add-support

Add support for environments and simplify migration syntax

src/houseplant/houseplant.py

@@ -14,6 +14,7 @@ class Houseplant:
     def __init__(self):
         self.console = Console()
         self.db = ClickHouseClient()
+        self.env = os.getenv("RAILS_ENV", "development")
 
     def init(self):
         """Initialize a new houseplant project."""
@@ -113,8 +114,10 @@ def migrate_up(self, version: str | None = None):
                 with open(os.path.join(migrations_dir, migration_file), "r") as f:
                     migration = yaml.safe_load(f)
 
-                if migration.get("up", {}).get("sql"):
-                    self.db.execute_migration(migration["up"]["sql"])
+                # Get migration SQL based on environment
+                migration_sql = migration.get(self.env, {}).get("up", {}).strip()
+                if migration_sql:
+                    self.db.execute_migration(migration_sql)
                     self.db.mark_migration_applied(migration_version)
                     self.console.print(
                         f"[green]✓[/green] Applied migration {migration_file}"
@@ -173,8 +176,10 @@ def migrate_down(self, version: str | None = None):
                 with open(os.path.join(migrations_dir, migration_file), "r") as f:
                     migration = yaml.safe_load(f)
 
-                if migration["down"]["sql"]:
-                    self.db.execute_migration(migration["down"]["sql"])
+                # Get migration SQL based on environment
+                migration_sql = migration.get(self.env, {}).get("down", {}).strip()
+                if migration_sql:
+                    self.db.execute_migration(migration_sql)
                     self.db.mark_migration_rolled_back(migration_version)
                     self.update_schema()
                     self.console.print(
@@ -201,11 +206,17 @@ def generate(self, name: str):
             with open(migration_file, "w") as f:
                 f.write(f"""version: "{timestamp}"
 name: {migration_name}
-up:
-  sql: ""
 
-down:
-  sql: ""
+development: &development
+  up: |
+  down: |
+
+test:
+  <<: *development
+
+production:
+  up: |
+  down: |
 """)
 
             self.console.print(f"✨ Generated migration: {migration_file}")

tests/test_houseplant.py

@@ -1,21 +1,116 @@
-#!/usr/bin/env python
-
-"""Tests for `houseplant` package."""
+import os
 
 import pytest
+import yaml
+
+from houseplant.houseplant import Houseplant
+
+
+@pytest.fixture
+def houseplant():
+    return Houseplant()
 
 
 @pytest.fixture
-def response():
-    """Sample pytest fixture.
+def test_migration(tmp_path):
+    # Set up test environment
+    migrations_dir = tmp_path / "ch/migrations"
+    migrations_dir.mkdir(parents=True)
+    migration_file = migrations_dir / "20240101000000_test_migration.yml"
+
+    migration_content = """version: "20240101000000"
+name: test_migration
+
+development: &development
+  up: |
+    CREATE TABLE dev_table (
+        id UInt32,
+        name String
+    ) ENGINE = MergeTree()
+    ORDER BY id
+  down: |
+    DROP TABLE dev_table
+
+test:
+  <<: *development
+
+production:
+  up: |
+    CREATE TABLE prod_table (
+        id UInt32,
+        name String
+    ) ENGINE = ReplicatedMergeTree('/clickhouse/{cluster}tables/{shard}/test/prod_table', '{replica}')
+    ORDER BY id
+  down: |
+    DROP TABLE prod_table
+"""
+
+    migration_file.write_text(migration_content)
+    os.chdir(tmp_path)
+    return migration_content
+
+
+def test_migrate_up_development(houseplant, test_migration, mocker):
+    # Mock environment and database calls
+    houseplant.env = "development"
+    mock_execute = mocker.patch.object(houseplant.db, "execute_migration")
+    mock_mark_applied = mocker.patch.object(houseplant.db, "mark_migration_applied")
+    mock_get_applied = mocker.patch.object(
+        houseplant.db, "get_applied_migrations", return_value=[]
+    )
+
+    # Run migration
+    houseplant.migrate_up()
+
+    # Verify correct SQL was executed
+    expected_sql = """CREATE TABLE dev_table (
+    id UInt32,
+    name String
+) ENGINE = MergeTree()
+ORDER BY id"""
+    mock_execute.assert_called_once_with(expected_sql)
+    mock_mark_applied.assert_called_once_with("20240101000000")
+    mock_get_applied.assert_called_once()
+
+
+def test_migrate_up_production(houseplant, test_migration, mocker):
+    # Mock environment and database calls
+    houseplant.env = "production"
+    mock_execute = mocker.patch.object(houseplant.db, "execute_migration")
+    mock_mark_applied = mocker.patch.object(houseplant.db, "mark_migration_applied")
+    mock_get_applied = mocker.patch.object(
+        houseplant.db, "get_applied_migrations", return_value=[]
+    )
+
+    # Run migration
+    houseplant.migrate_up()
+
+    # Verify correct SQL was executed
+    expected_sql = """CREATE TABLE prod_table (
+    id UInt32,
+    name String
+) ENGINE = ReplicatedMergeTree('/clickhouse/{cluster}tables/{shard}/test/prod_table', '{replica}')
+ORDER BY id"""
+    mock_execute.assert_called_once_with(expected_sql)
+    mock_mark_applied.assert_called_once_with("20240101000000")
+    mock_get_applied.assert_called_once()
+
 
-    See more at: http://doc.pytest.org/en/latest/fixture.html
-    """
-    # import requests
-    # return requests.get('https://github.com/audreyr/cookiecutter-pypackage')
+@pytest.mark.skip
+def test_migrate_down(houseplant, test_migration, mocker):
+    # Mock database calls
+    mock_execute = mocker.patch.object(houseplant.db, "execute_migration")
+    mock_mark_rolled_back = mocker.patch.object(
+        houseplant.db, "mark_migration_rolled_back"
+    )
+    mock_get_applied = mocker.patch.object(
+        houseplant.db, "get_applied_migrations", return_value=[("20240101000000",)]
+    )
 
+    # Roll back migration
+    houseplant.migrate_down()
 
-def test_content(response):
-    """Sample pytest test function with the pytest fixture as an argument."""
-    # from bs4 import BeautifulSoup
-    # assert 'GitHub' in BeautifulSoup(response.content).title.string
+    # Verify correct SQL was executed
+    mock_execute.assert_called_once_with("DROP TABLE dev_table")
+    mock_mark_rolled_back.assert_called_once_with("20240101000000")
+    mock_get_applied.assert_called_once()